From 34807188cc4454d4e951cc1b304db10d025e5824 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 25 Feb 2025 13:18:08 -0800 Subject: [PATCH 001/122] initial implementation for pytest benchmark discovery --- .../tests/pytest/test_benchmark_bubble_sort.py | 6 ++++++ codeflash/discovery/discover_unit_tests.py | 2 ++ .../discovery/pytest_new_process_discovery.py | 12 +++++++++++- codeflash/verification/test_results.py | 10 +++++++--- tests/test_unit_test_discovery.py | 14 ++++++++++++++ 5 files changed, 40 insertions(+), 4 deletions(-) create mode 100644 code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py diff --git a/code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py new file mode 100644 index 000000000..dcbb86ac1 --- /dev/null +++ b/code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py @@ -0,0 +1,6 @@ +from code_to_optimize.bubble_sort import sorter + + +def test_sort(benchmark): + result = benchmark(sorter, list(reversed(range(5000)))) + assert result == list(range(5000)) diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py index a36774f66..a4b0f8ac9 100644 --- a/codeflash/discovery/discover_unit_tests.py +++ b/codeflash/discovery/discover_unit_tests.py @@ -97,6 +97,8 @@ def discover_tests_pytest( test_type = TestType.REPLAY_TEST elif "test_concolic_coverage" in test["test_file"]: test_type = TestType.CONCOLIC_COVERAGE_TEST + elif test["test_type"] == "benchmark": # New condition for benchmark tests + test_type = TestType.BENCHMARK_TEST else: test_type = TestType.EXISTING_UNIT_TEST diff --git a/codeflash/discovery/pytest_new_process_discovery.py b/codeflash/discovery/pytest_new_process_discovery.py index e85bbdcce..83175218b 100644 --- a/codeflash/discovery/pytest_new_process_discovery.py +++ b/codeflash/discovery/pytest_new_process_discovery.py @@ -23,7 +23,17 @@ def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, s test_class = None if test.cls: test_class = test.parent.name - test_results.append({"test_file": str(test.path), "test_class": test_class, "test_function": test.name}) + + # Determine if this is a benchmark test by checking for the benchmark fixture + is_benchmark = hasattr(test, 'fixturenames') and 'benchmark' in test.fixturenames + test_type = 'benchmark' if is_benchmark else 'regular' + + test_results.append({ + "test_file": str(test.path), + "test_class": test_class, + "test_function": test.name, + "test_type": test_type + }) return test_results diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py index 28d8bfc0d..99151f983 100644 --- a/codeflash/verification/test_results.py +++ b/codeflash/verification/test_results.py @@ -29,6 +29,7 @@ class TestType(Enum): REPLAY_TEST = 4 CONCOLIC_COVERAGE_TEST = 5 INIT_STATE_TEST = 6 + BENCHMARK_TEST = 7 def to_name(self) -> str: if self == TestType.INIT_STATE_TEST: @@ -39,6 +40,7 @@ def to_name(self) -> str: TestType.GENERATED_REGRESSION: "πŸŒ€ Generated Regression Tests", TestType.REPLAY_TEST: "βͺ Replay Tests", TestType.CONCOLIC_COVERAGE_TEST: "πŸ”Ž Concolic Coverage Tests", + TestType.BENCHMARK_TEST: "πŸ“ Benchmark Tests", } return names[self] @@ -66,6 +68,7 @@ def from_str_id(string_id: str, iteration_id: Optional[str] = None) -> Invocatio else: test_class_name = second_components[0] test_function_name = second_components[1] + # logger.debug(f"Invocation id info: test_module_path: {components[0]}, test_class_name: {test_class_name}, test_function_name: {test_function_name}, function_getting_tested: {components[2]}, iteration_id: {iteration_id if iteration_id else components[3]}") return InvocationId( test_module_path=components[0], test_class_name=test_class_name, @@ -167,9 +170,10 @@ def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree: def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: for result in self.test_results: if result.did_pass and not result.runtime: - logger.debug( - f"Ignoring test case that passed but had no runtime -> {result.id}, Loop # {result.loop_index}, Test Type: {result.test_type}, Verification Type: {result.verification_type}" - ) + pass + # logger.debug( + # f"Ignoring test case that passed but had no runtime -> {result.id}, Loop # {result.loop_index}, Test Type: {result.test_type}, Verification Type: {result.verification_type}" + # ) usable_runtimes = [ (result.id, result.runtime) for result in self.test_results if result.did_pass and result.runtime ] diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index 3e5bfa120..b8d86b70c 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -19,6 +19,20 @@ def test_unit_test_discovery_pytest(): assert len(tests) > 0 # print(tests) +def test_benchmark_test_discovery_pytest(): + project_path = Path(__file__).parent.parent.resolve() / "code_to_optimize" + tests_path = project_path / "tests" / "pytest" + test_config = TestConfig( + tests_root=tests_path, + project_root_path=project_path, + test_framework="pytest", + tests_project_rootdir=tests_path.parent, + ) + tests = discover_unit_tests(test_config) + print(tests) + assert len(tests) > 0 + # print(tests) + def test_unit_test_discovery_unittest(): project_path = Path(__file__).parent.parent.resolve() / "code_to_optimize" From 133a9e3635c64a9a6d4739e146ca5e1f18b89fea Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 27 Feb 2025 16:25:07 -0800 Subject: [PATCH 002/122] initial implementation for tracing benchmarks using a plugin, and projecting speedup --- code_to_optimize/process_and_bubble_sort.py | 28 +++++ .../test_benchmark_bubble_sort.py | 7 ++ .../benchmarks/test_process_and_sort.py | 8 ++ codeflash/benchmarking/__init__.py | 0 codeflash/benchmarking/get_trace_info.py | 112 ++++++++++++++++++ codeflash/benchmarking/plugin/__init__.py | 0 codeflash/benchmarking/plugin/plugin.py | 79 ++++++++++++ .../pytest_new_process_trace_benchmarks.py | 15 +++ codeflash/benchmarking/trace_benchmarks.py | 20 ++++ codeflash/cli_cmds/cli.py | 9 +- codeflash/discovery/discover_unit_tests.py | 2 - codeflash/discovery/functions_to_optimize.py | 6 +- .../pytest_new_process_discover_benchmarks.py | 54 +++++++++ .../discovery/pytest_new_process_discovery.py | 12 +- codeflash/optimization/function_optimizer.py | 21 +++- codeflash/optimization/optimizer.py | 37 +++++- codeflash/tracer.py | 52 +++++--- codeflash/verification/test_results.py | 2 - codeflash/verification/verification_utils.py | 1 + pyproject.toml | 12 +- tests/test_trace_benchmarks.py | 8 ++ tests/test_unit_test_discovery.py | 8 +- 22 files changed, 446 insertions(+), 47 deletions(-) create mode 100644 code_to_optimize/process_and_bubble_sort.py rename code_to_optimize/tests/pytest/{ => benchmarks}/test_benchmark_bubble_sort.py (50%) create mode 100644 code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py create mode 100644 codeflash/benchmarking/__init__.py create mode 100644 codeflash/benchmarking/get_trace_info.py create mode 100644 codeflash/benchmarking/plugin/__init__.py create mode 100644 codeflash/benchmarking/plugin/plugin.py create mode 100644 codeflash/benchmarking/pytest_new_process_trace_benchmarks.py create mode 100644 codeflash/benchmarking/trace_benchmarks.py create mode 100644 codeflash/discovery/pytest_new_process_discover_benchmarks.py create mode 100644 tests/test_trace_benchmarks.py diff --git a/code_to_optimize/process_and_bubble_sort.py b/code_to_optimize/process_and_bubble_sort.py new file mode 100644 index 000000000..94359e599 --- /dev/null +++ b/code_to_optimize/process_and_bubble_sort.py @@ -0,0 +1,28 @@ +from code_to_optimize.bubble_sort import sorter + + +def calculate_pairwise_products(arr): + """ + Calculate the average of all pairwise products in the array. + """ + sum_of_products = 0 + count = 0 + + for i in range(len(arr)): + for j in range(len(arr)): + if i != j: + sum_of_products += arr[i] * arr[j] + count += 1 + + # The average of all pairwise products + return sum_of_products / count if count > 0 else 0 + + +def compute_and_sort(arr): + # Compute pairwise sums average + pairwise_average = calculate_pairwise_products(arr) + + # Call sorter function + sorter(arr.copy()) + + return pairwise_average diff --git a/code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py similarity index 50% rename from code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py rename to code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py index dcbb86ac1..f1ebcf5c7 100644 --- a/code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py @@ -1,6 +1,13 @@ +import pytest + from code_to_optimize.bubble_sort import sorter def test_sort(benchmark): result = benchmark(sorter, list(reversed(range(5000)))) assert result == list(range(5000)) + +# This should not be picked up as a benchmark test +def test_sort2(): + result = sorter(list(reversed(range(5000)))) + assert result == list(range(5000)) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py new file mode 100644 index 000000000..ca2f0ef65 --- /dev/null +++ b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py @@ -0,0 +1,8 @@ +from code_to_optimize.process_and_bubble_sort import compute_and_sort +from code_to_optimize.bubble_sort2 import sorter +def test_compute_and_sort(benchmark): + result = benchmark(compute_and_sort, list(reversed(range(5000)))) + assert result == 6247083.5 + +def test_no_func(benchmark): + benchmark(sorter, list(reversed(range(5000)))) \ No newline at end of file diff --git a/codeflash/benchmarking/__init__.py b/codeflash/benchmarking/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/codeflash/benchmarking/get_trace_info.py b/codeflash/benchmarking/get_trace_info.py new file mode 100644 index 000000000..1d0b339d9 --- /dev/null +++ b/codeflash/benchmarking/get_trace_info.py @@ -0,0 +1,112 @@ +import sqlite3 +from pathlib import Path +from typing import Dict, Set + +from codeflash.discovery.functions_to_optimize import FunctionToOptimize + + +def get_function_benchmark_timings(trace_dir: Path, all_functions_to_optimize: list[FunctionToOptimize]) -> dict[str, dict[str, float]]: + """Process all trace files in the given directory and extract timing data for the specified functions. + + Args: + trace_dir: Path to the directory containing .trace files + all_functions_to_optimize: Set of FunctionToOptimize objects representing functions to include + + Returns: + A nested dictionary where: + - Outer keys are function qualified names with file name + - Inner keys are benchmark names (trace filename without .trace extension) + - Values are function timing in milliseconds + + """ + # Create a mapping of (filename, function_name, class_name) -> qualified_name for efficient lookups + function_lookup = {} + function_benchmark_timings = {} + + for func in all_functions_to_optimize: + qualified_name = func.qualified_name_with_file_name + + # Extract components (assumes Path.name gives only filename without directory) + filename = func.file_path + function_name = func.function_name + + # Get class name if there's a parent + class_name = func.parents[0].name if func.parents else None + + # Store in lookup dictionary + key = (filename, function_name, class_name) + function_lookup[key] = qualified_name + function_benchmark_timings[qualified_name] = {} + + # Find all .trace files in the directory + trace_files = list(trace_dir.glob("*.trace")) + + for trace_file in trace_files: + # Extract benchmark name from filename (without .trace) + benchmark_name = trace_file.stem + + # Connect to the trace database + conn = sqlite3.connect(trace_file) + cursor = conn.cursor() + + # For each function we're interested in, query the database directly + for (filename, function_name, class_name), qualified_name in function_lookup.items(): + # Adjust query based on whether we have a class name + if class_name: + cursor.execute( + "SELECT total_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND class_name = ?", + (f"%{filename}", function_name, class_name) + ) + else: + cursor.execute( + "SELECT total_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND (class_name IS NULL OR class_name = '')", + (f"%{filename}", function_name) + ) + + result = cursor.fetchone() + if result: + time_ns = result[0] + function_benchmark_timings[qualified_name][benchmark_name] = time_ns / 1e6 # Convert to milliseconds + + conn.close() + + return function_benchmark_timings + + +def get_benchmark_timings(trace_dir: Path) -> dict[str, float]: + """Extract total benchmark timings from trace files. + + Args: + trace_dir: Path to the directory containing .trace files + + Returns: + A dictionary mapping benchmark names to their total execution time in milliseconds. + """ + benchmark_timings = {} + + # Find all .trace files in the directory + trace_files = list(trace_dir.glob("*.trace")) + + for trace_file in trace_files: + # Extract benchmark name from filename (without .trace extension) + benchmark_name = trace_file.stem + + # Connect to the trace database + conn = sqlite3.connect(trace_file) + cursor = conn.cursor() + + # Query the total_time table for the benchmark's total execution time + try: + cursor.execute("SELECT time_ns FROM total_time") + result = cursor.fetchone() + if result: + time_ns = result[0] + # Convert nanoseconds to milliseconds + benchmark_timings[benchmark_name] = time_ns / 1e6 + except sqlite3.OperationalError: + # Handle case where total_time table might not exist + print(f"Warning: Could not get total time for benchmark {benchmark_name}") + + conn.close() + + return benchmark_timings diff --git a/codeflash/benchmarking/plugin/__init__.py b/codeflash/benchmarking/plugin/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py new file mode 100644 index 000000000..34ca2b777 --- /dev/null +++ b/codeflash/benchmarking/plugin/plugin.py @@ -0,0 +1,79 @@ +import pytest + +from codeflash.tracer import Tracer +from pathlib import Path + +class CodeFlashPlugin: + @staticmethod + def pytest_addoption(parser): + parser.addoption( + "--codeflash-trace", + action="store_true", + 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): + if hasattr(plugin, "name") and plugin.name == "pytest-benchmark": + manager.unregister(plugin) + + @staticmethod + def pytest_collection_modifyitems(config, items): + 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) + + @staticmethod + @pytest.fixture + def benchmark(request): + if not request.config.getoption("--codeflash-trace"): + return None + + 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 + ) + + with tracer: + result = func(*args, **kwargs) + + return result + + return Benchmark() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py new file mode 100644 index 000000000..b892d62a0 --- /dev/null +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -0,0 +1,15 @@ +import sys +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], plugins=[CodeFlashPlugin()] + ) + except Exception as e: + print(f"Failed to collect tests: {e!s}") + exitcode = -1 \ No newline at end of file diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py new file mode 100644 index 000000000..2d3acdd66 --- /dev/null +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -0,0 +1,20 @@ +from __future__ import annotations +from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE +from pathlib import Path +import subprocess + +def trace_benchmarks_pytest(benchmarks_root: Path, project_root: Path, function_list: list[str] = []) -> None: + result = subprocess.run( + [ + SAFE_SYS_EXECUTABLE, + Path(__file__).parent / "pytest_new_process_trace_benchmarks.py", + str(benchmarks_root), + ",".join(function_list) + ], + cwd=project_root, + check=False, + capture_output=True, + text=True, + ) + print("stdout:", result.stdout) + print("stderr:", result.stderr) diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 012fc86eb..04445f1db 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -62,6 +62,10 @@ def parse_args() -> Namespace: ) parser.add_argument("-v", "--verbose", action="store_true", help="Print verbose debug logs") parser.add_argument("--version", action="store_true", help="Print the version of codeflash") + parser.add_argument("--benchmark", action="store_true", help="Trace benchmark tests and calculate optimization impact on benchmarks") + parser.add_argument( + "--benchmarks-root", type=str, help="Path to the directory of the project, where all the pytest-benchmark tests are located." + ) args: Namespace = parser.parse_args() return process_and_validate_cmd_args(args) @@ -116,6 +120,7 @@ def process_pyproject_config(args: Namespace) -> Namespace: "disable_telemetry", "disable_imports_sorting", "git_remote", + "benchmarks_root" ] for key in supported_keys: if key in pyproject_config and ( @@ -127,7 +132,9 @@ def process_pyproject_config(args: Namespace) -> Namespace: assert Path(args.module_root).is_dir(), f"--module-root {args.module_root} must be a valid directory" assert args.tests_root is not None, "--tests-root must be specified" assert Path(args.tests_root).is_dir(), f"--tests-root {args.tests_root} must be a valid directory" - + if args.benchmark: + assert args.benchmarks_root is not None, "--benchmarks-root must be specified when running with --benchmark" + assert Path(args.benchmarks_root).is_dir(), f"--benchmarks-root {args.benchmarks_root} must be a valid directory" assert not (env_utils.get_pr_number() is not None and not env_utils.ensure_codeflash_api_key()), ( "Codeflash API key not found. When running in a Github Actions Context, provide the " "'CODEFLASH_API_KEY' environment variable as a secret.\n" diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py index a4b0f8ac9..a36774f66 100644 --- a/codeflash/discovery/discover_unit_tests.py +++ b/codeflash/discovery/discover_unit_tests.py @@ -97,8 +97,6 @@ def discover_tests_pytest( test_type = TestType.REPLAY_TEST elif "test_concolic_coverage" in test["test_file"]: test_type = TestType.CONCOLIC_COVERAGE_TEST - elif test["test_type"] == "benchmark": # New condition for benchmark tests - test_type = TestType.BENCHMARK_TEST else: test_type = TestType.EXISTING_UNIT_TEST diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index 11eebab16..71c63ddc4 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -121,7 +121,6 @@ class FunctionToOptimize: method extends this with the module name from the project root. """ - function_name: str file_path: Path parents: list[FunctionParent] # list[ClassDef | FunctionDef | AsyncFunctionDef] @@ -145,6 +144,11 @@ def qualified_name(self) -> str: def qualified_name_with_modules_from_root(self, project_root_path: Path) -> str: return f"{module_name_from_file_path(self.file_path, project_root_path)}.{self.qualified_name}" + @property + def qualified_name_with_file_name(self) -> str: + class_name = self.parents[0].name if self.parents else None + return f"{self.file_path}:{(class_name + ':' if class_name else '')}{self.function_name}" + def get_functions_to_optimize( optimize_all: str | None, diff --git a/codeflash/discovery/pytest_new_process_discover_benchmarks.py b/codeflash/discovery/pytest_new_process_discover_benchmarks.py new file mode 100644 index 000000000..83175218b --- /dev/null +++ b/codeflash/discovery/pytest_new_process_discover_benchmarks.py @@ -0,0 +1,54 @@ +import sys +from typing import Any + +# This script should not have any relation to the codeflash package, be careful with imports +cwd = sys.argv[1] +tests_root = sys.argv[2] +pickle_path = sys.argv[3] +collected_tests = [] +pytest_rootdir = None +sys.path.insert(1, str(cwd)) + + +class PytestCollectionPlugin: + def pytest_collection_finish(self, session) -> None: + global pytest_rootdir + collected_tests.extend(session.items) + pytest_rootdir = session.config.rootdir + + +def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, str]]: + test_results = [] + for test in pytest_tests: + test_class = None + if test.cls: + test_class = test.parent.name + + # Determine if this is a benchmark test by checking for the benchmark fixture + is_benchmark = hasattr(test, 'fixturenames') and 'benchmark' in test.fixturenames + test_type = 'benchmark' if is_benchmark else 'regular' + + test_results.append({ + "test_file": str(test.path), + "test_class": test_class, + "test_function": test.name, + "test_type": test_type + }) + return test_results + + +if __name__ == "__main__": + import pytest + + try: + exitcode = pytest.main( + [tests_root, "-pno:logging", "--collect-only", "-m", "not skip"], plugins=[PytestCollectionPlugin()] + ) + except Exception as e: + print(f"Failed to collect tests: {e!s}") + exitcode = -1 + tests = parse_pytest_collection_results(collected_tests) + import pickle + + with open(pickle_path, "wb") as f: + pickle.dump((exitcode, tests, pytest_rootdir), f, protocol=pickle.HIGHEST_PROTOCOL) diff --git a/codeflash/discovery/pytest_new_process_discovery.py b/codeflash/discovery/pytest_new_process_discovery.py index 83175218b..e85bbdcce 100644 --- a/codeflash/discovery/pytest_new_process_discovery.py +++ b/codeflash/discovery/pytest_new_process_discovery.py @@ -23,17 +23,7 @@ def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, s test_class = None if test.cls: test_class = test.parent.name - - # Determine if this is a benchmark test by checking for the benchmark fixture - is_benchmark = hasattr(test, 'fixturenames') and 'benchmark' in test.fixturenames - test_type = 'benchmark' if is_benchmark else 'regular' - - test_results.append({ - "test_file": str(test.path), - "test_class": test_class, - "test_function": test.name, - "test_type": test_type - }) + test_results.append({"test_file": str(test.path), "test_class": test_class, "test_function": test.name}) return test_results diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 7b067a094..66d3c6ab6 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -92,6 +92,8 @@ def __init__( function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_ast: ast.FunctionDef | None = None, aiservice_client: AiServiceClient | None = None, + function_benchmark_timings: dict[str, dict[str, float]] | None = None, + total_benchmark_timings: dict[str, float] | None = None, args: Namespace | None = None, ) -> None: self.project_root = test_cfg.project_root_path @@ -120,6 +122,9 @@ def __init__( self.function_trace_id: str = str(uuid.uuid4()) self.original_module_path = module_name_from_file_path(self.function_to_optimize.file_path, self.project_root) + self.function_benchmark_timings = function_benchmark_timings if function_benchmark_timings else {} + self.total_benchmark_timings = total_benchmark_timings if total_benchmark_timings else {} + def optimize_function(self) -> Result[BestOptimization, str]: should_run_experiment = self.experiment_id is not None logger.debug(f"Function Trace ID: {self.function_trace_id}") @@ -280,6 +285,20 @@ def optimize_function(self) -> Result[BestOptimization, str]: function_name=function_to_optimize_qualified_name, file_path=self.function_to_optimize.file_path, ) + speedup = explanation.speedup # eg. 1.2 means 1.2x faster + if self.args.benchmark: + fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_file_name] + for benchmark_name, og_benchmark_timing in fto_benchmark_timings.items(): + print(f"Calculating speedup for benchmark {benchmark_name}") + total_benchmark_timing = self.total_benchmark_timings[benchmark_name] + # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values + expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + og_benchmark_timing / speedup + print(f"Expected new benchmark timing: {expected_new_benchmark_timing}") + print(f"Original benchmark timing: {total_benchmark_timing}") + print(f"Benchmark speedup: {total_benchmark_timing / expected_new_benchmark_timing}") + + speedup = total_benchmark_timing / expected_new_benchmark_timing + print(f"Speedup: {speedup}") self.log_successful_optimization(explanation, generated_tests) @@ -1107,7 +1126,7 @@ def run_and_parse_tests( f"stdout: {run_result.stdout}\n" f"stderr: {run_result.stderr}\n" ) - + # print(test_files) results, coverage_results = parse_test_results( test_xml_path=result_file_path, test_files=test_files, diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index cae78a153..01a196143 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient +from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.cli_cmds.console import console, logger from codeflash.code_utils import env_utils from codeflash.code_utils.code_replacer import normalize_code, normalize_node @@ -21,6 +22,7 @@ from codeflash.telemetry.posthog_cf import ph from codeflash.verification.test_results import TestType from codeflash.verification.verification_utils import TestConfig +from codeflash.benchmarking.get_trace_info import get_function_benchmark_timings, get_benchmark_timings if TYPE_CHECKING: from argparse import Namespace @@ -51,6 +53,8 @@ def create_function_optimizer( function_to_optimize_ast: ast.FunctionDef | None = None, function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_source_code: str | None = "", + function_benchmark_timings: dict[str, dict[str, float]] | None = None, + total_benchmark_timings: dict[str, float] | None = None, ) -> FunctionOptimizer: return FunctionOptimizer( function_to_optimize=function_to_optimize, @@ -60,7 +64,8 @@ def create_function_optimizer( function_to_optimize_ast=function_to_optimize_ast, aiservice_client=self.aiservice_client, args=self.args, - + function_benchmark_timings=function_benchmark_timings if function_benchmark_timings else None, + total_benchmark_timings=total_benchmark_timings if total_benchmark_timings else None, ) def run(self) -> None: @@ -82,6 +87,23 @@ def run(self) -> None: project_root=self.args.project_root, 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) + print(function_benchmark_timings) + total_benchmark_timings = get_benchmark_timings(trace_dir) + print("Total benchmark timings:") + print(total_benchmark_timings) + # for function in fully_qualified_function_names: + optimizations_found: int = 0 function_iterator_count: int = 0 @@ -160,10 +182,17 @@ def run(self) -> None: f"Skipping optimization." ) continue + if self.args.benchmark: + + function_optimizer = self.create_function_optimizer( + function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code, function_benchmark_timings, total_benchmark_timings + ) + else: + function_optimizer = self.create_function_optimizer( + function_to_optimize, function_to_optimize_ast, function_to_tests, + validated_original_code[original_module_path].source_code + ) - function_optimizer = self.create_function_optimizer( - function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code - ) best_optimization = function_optimizer.optimize_function() if is_successful(best_optimization): optimizations_found += 1 diff --git a/codeflash/tracer.py b/codeflash/tracer.py index 39b05e01f..5bc1ae482 100644 --- a/codeflash/tracer.py +++ b/codeflash/tracer.py @@ -40,7 +40,9 @@ from codeflash.tracing.replay_test import create_trace_replay_test from codeflash.tracing.tracing_utils import FunctionModules from codeflash.verification.verification_utils import get_test_file_path - +# import warnings +# warnings.filterwarnings("ignore", category=dill.PickleWarning) +# warnings.filterwarnings("ignore", category=DeprecationWarning) # Debug this file by simply adding print statements. This file is not meant to be debugged by the debugger. class Tracer: @@ -117,14 +119,15 @@ def __init__( def __enter__(self) -> None: if self.disable: return - if getattr(Tracer, "used_once", False): - console.print( - "Codeflash: Tracer can only be used once per program run. " - "Please only enable the Tracer once. Skipping tracing this section." - ) - self.disable = True - return - Tracer.used_once = True + + # if getattr(Tracer, "used_once", False): + # console.print( + # "Codeflash: Tracer can only be used once per program run. " + # "Please only enable the Tracer once. Skipping tracing this section." + # ) + # self.disable = True + # return + # Tracer.used_once = True if pathlib.Path(self.output_file).exists(): console.print("Codeflash: Removing existing trace file") @@ -149,6 +152,14 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: return sys.setprofile(None) self.con.commit() + # Check if any functions were actually traced + if self.trace_count == 0: + self.con.close() + # Delete the trace file if no functions were traced + if self.output_file.exists(): + self.output_file.unlink() + console.print("Codeflash: No functions were traced. Removing trace database.") + return self.create_stats() @@ -193,7 +204,9 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: test_framework=self.config["test_framework"], max_run_count=self.max_function_count, ) - function_path = "_".join(self.functions) if self.functions else self.file_being_called_from + # Need a better way to store the replay test + # function_path = "_".join(self.functions) if self.functions else self.file_being_called_from + function_path = self.file_being_called_from test_file_path = get_test_file_path( test_dir=Path(self.config["tests_root"]), function_name=function_path, test_type="replay" ) @@ -224,9 +237,9 @@ def tracer_logic(self, frame: FrameType, event: str): return if not file_name.exists(): return - if self.functions: - if code.co_name not in self.functions: - return + # if self.functions: + # if code.co_name not in self.functions: + # return class_name = None arguments = frame.f_locals try: @@ -241,9 +254,13 @@ def tracer_logic(self, frame: FrameType, event: str): except: # someone can override the getattr method and raise an exception. I'm looking at you wrapt return + function_qualified_name = f"{file_name}:{(class_name + ':' if class_name else '')}{code.co_name}" if function_qualified_name in self.ignored_qualified_functions: return + if self.functions and function_qualified_name not in self.functions: + return + if function_qualified_name not in self.function_count: # seeing this function for the first time self.function_count[function_qualified_name] = 0 @@ -476,7 +493,7 @@ def print_stats(self, sort=-1): # print in milliseconds. s = StringIO() stats_obj = pstats.Stats(copy(self), stream=s) - stats_obj.strip_dirs().sort_stats(*sort).print_stats(25) + stats_obj.strip_dirs().sort_stats(*sort).print_stats(100) self.total_tt = stats_obj.total_tt console.print("total_tt", self.total_tt) raw_stats = s.getvalue() @@ -621,13 +638,16 @@ def main(): "__cached__": None, } try: - Tracer( + tracer = Tracer( output=args.outfile, functions=args.only_functions, max_function_count=args.max_function_count, timeout=args.tracer_timeout, config_file_path=args.codeflash_config, - ).runctx(code, globs, None) + ) + + tracer.runctx(code, globs, None) + print(tracer.functions) except BrokenPipeError as exc: # Prevent "Exception ignored" during interpreter shutdown. diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py index 99151f983..db01ff049 100644 --- a/codeflash/verification/test_results.py +++ b/codeflash/verification/test_results.py @@ -29,7 +29,6 @@ class TestType(Enum): REPLAY_TEST = 4 CONCOLIC_COVERAGE_TEST = 5 INIT_STATE_TEST = 6 - BENCHMARK_TEST = 7 def to_name(self) -> str: if self == TestType.INIT_STATE_TEST: @@ -40,7 +39,6 @@ def to_name(self) -> str: TestType.GENERATED_REGRESSION: "πŸŒ€ Generated Regression Tests", TestType.REPLAY_TEST: "βͺ Replay Tests", TestType.CONCOLIC_COVERAGE_TEST: "πŸ”Ž Concolic Coverage Tests", - TestType.BENCHMARK_TEST: "πŸ“ Benchmark Tests", } return names[self] diff --git a/codeflash/verification/verification_utils.py b/codeflash/verification/verification_utils.py index 3d30f89f9..fdff3c935 100644 --- a/codeflash/verification/verification_utils.py +++ b/codeflash/verification/verification_utils.py @@ -74,3 +74,4 @@ class TestConfig: # or for unittest - project_root_from_module_root(args.tests_root, pyproject_file_path) concolic_test_root_dir: Optional[Path] = None pytest_cmd: str = "pytest" + benchmark_tests_root: Optional[Path] = None diff --git a/pyproject.toml b/pyproject.toml index 27ebf6c1b..026c0eafd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -211,13 +211,13 @@ initial-content = """ [tool.codeflash] -module-root = "codeflash" -tests-root = "tests" +# All paths are relative to this pyproject.toml's directory. +module-root = "code_to_optimize" +tests-root = "code_to_optimize/tests" +benchmarks-root = "code_to_optimize/tests/pytest/benchmarks" test-framework = "pytest" -formatter-cmds = [ - "poetry run ruff check --exit-zero --fix $file", - "poetry run ruff format $file", -] +ignore-paths = [] +formatter-cmds = ["ruff check --exit-zero --fix $file", "ruff format $file"] [build-system] diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py new file mode 100644 index 000000000..acd40b0b3 --- /dev/null +++ b/tests/test_trace_benchmarks.py @@ -0,0 +1,8 @@ +from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest +from pathlib import Path + +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"]) \ No newline at end of file diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index b8d86b70c..fe56b907f 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -3,6 +3,7 @@ from pathlib import Path from codeflash.discovery.discover_unit_tests import discover_unit_tests +from codeflash.verification.test_results import TestType from codeflash.verification.verification_utils import TestConfig @@ -21,7 +22,7 @@ def test_unit_test_discovery_pytest(): def test_benchmark_test_discovery_pytest(): project_path = Path(__file__).parent.parent.resolve() / "code_to_optimize" - tests_path = project_path / "tests" / "pytest" + tests_path = project_path / "tests" / "pytest" / "benchmarks" / "test_benchmark_bubble_sort.py" test_config = TestConfig( tests_root=tests_path, project_root_path=project_path, @@ -29,9 +30,10 @@ def test_benchmark_test_discovery_pytest(): tests_project_rootdir=tests_path.parent, ) tests = discover_unit_tests(test_config) - print(tests) assert len(tests) > 0 - # print(tests) + assert 'bubble_sort.sorter' in tests + benchmark_tests = sum(1 for test in tests['bubble_sort.sorter'] if test.tests_in_file.test_type == TestType.BENCHMARK_TEST) + assert benchmark_tests == 1 def test_unit_test_discovery_unittest(): From 2f26695f20a45128e44ba05f3d605b300a9c5473 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 4 Mar 2025 15:53:24 -0800 Subject: [PATCH 003/122] initial implementation of tracing benchmarks via the plugin --- codeflash/benchmarking/get_trace_info.py | 10 ++-- codeflash/benchmarking/plugin/plugin.py | 3 +- .../pytest_new_process_trace_benchmarks.py | 2 +- codeflash/benchmarking/utils.py | 26 ++++++++++ codeflash/cli_cmds/cli.py | 1 + codeflash/code_utils/config_parser.py | 4 +- codeflash/models/models.py | 4 +- codeflash/optimization/function_optimizer.py | 51 +++++++++++++++---- codeflash/optimization/optimizer.py | 13 +++-- codeflash/tracer.py | 40 ++++++++++----- codeflash/verification/test_results.py | 29 +++++++++++ codeflash/verification/test_runner.py | 2 + 12 files changed, 149 insertions(+), 36 deletions(-) create mode 100644 codeflash/benchmarking/utils.py diff --git a/codeflash/benchmarking/get_trace_info.py b/codeflash/benchmarking/get_trace_info.py index 1d0b339d9..3dd3831ce 100644 --- a/codeflash/benchmarking/get_trace_info.py +++ b/codeflash/benchmarking/get_trace_info.py @@ -54,18 +54,20 @@ def get_function_benchmark_timings(trace_dir: Path, all_functions_to_optimize: l # Adjust query based on whether we have a class name if class_name: cursor.execute( - "SELECT total_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND class_name = ?", + "SELECT cumulative_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND class_name = ?", (f"%{filename}", function_name, class_name) ) else: cursor.execute( - "SELECT total_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND (class_name IS NULL OR class_name = '')", + "SELECT cumulative_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND (class_name IS NULL OR class_name = '')", (f"%{filename}", function_name) ) - result = cursor.fetchone() + result = cursor.fetchall() + if len(result) > 1: + print(f"Multiple results found for {qualified_name} in {benchmark_name}: {result}") if result: - time_ns = result[0] + time_ns = result[0][0] function_benchmark_timings[qualified_name][benchmark_name] = time_ns / 1e6 # Convert to milliseconds conn.close() diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index 34ca2b777..80accec22 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -68,7 +68,8 @@ def __call__(self, func, *args, **kwargs): tracer = Tracer( output=str(output_path), # Convert Path to string for Tracer functions=trace_functions, - max_function_count=256 + max_function_count=256, + benchmark=True ) with tracer: diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index b892d62a0..6b91e2b4f 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -8,7 +8,7 @@ try: exitcode = pytest.main( - [benchmarks_root, "--benchmarks-root", benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "--functions", function_list], plugins=[CodeFlashPlugin()] + [benchmarks_root, "--benchmarks-root", benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "--functions", function_list,"-o", "addopts="], plugins=[CodeFlashPlugin()] ) except Exception as e: print(f"Failed to collect tests: {e!s}") diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py new file mode 100644 index 000000000..d97c2e36e --- /dev/null +++ b/codeflash/benchmarking/utils.py @@ -0,0 +1,26 @@ +def print_benchmark_table(function_benchmark_timings, total_benchmark_timings): + # Print table header + print(f"{'Benchmark Test':<50} | {'Total Time (s)':<15} | {'Function Time (s)':<15} | {'Percentage (%)':<15}") + print("-" * 100) + + # Process each function's benchmark data + for func_path, test_times in function_benchmark_timings.items(): + function_name = func_path.split(":")[-1] + print(f"\n== Function: {function_name} ==") + + # Sort by percentage (highest first) + sorted_tests = [] + for test_name, func_time in test_times.items(): + total_time = total_benchmark_timings.get(test_name, 0) + if total_time > 0: + percentage = (func_time / total_time) * 100 + sorted_tests.append((test_name, total_time, func_time, percentage)) + + sorted_tests.sort(key=lambda x: x[3], reverse=True) + + # Print each test's data + for test_name, total_time, func_time, percentage in sorted_tests: + print(f"{test_name:<50} | {total_time:<15.3f} | {func_time:<15.3f} | {percentage:<15.2f}") + +# Usage + diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 04445f1db..96bb0cef3 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -113,6 +113,7 @@ def process_pyproject_config(args: Namespace) -> Namespace: supported_keys = [ "module_root", "tests_root", + "benchmarks_root", "test_framework", "ignore_paths", "pytest_cmd", diff --git a/codeflash/code_utils/config_parser.py b/codeflash/code_utils/config_parser.py index 89832503b..1edd24fa8 100644 --- a/codeflash/code_utils/config_parser.py +++ b/codeflash/code_utils/config_parser.py @@ -50,10 +50,10 @@ def parse_config_file(config_file_path: Path | None = None) -> tuple[dict[str, A assert isinstance(config, dict) # default values: - path_keys = ["module-root", "tests-root"] + path_keys = ["module-root", "tests-root", "benchmarks-root"] path_list_keys = ["ignore-paths"] str_keys = {"pytest-cmd": "pytest", "git-remote": "origin"} - bool_keys = {"disable-telemetry": False, "disable-imports-sorting": False} + bool_keys = {"disable-telemetry": False, "disable-imports-sorting": False, "benchmark": False} list_str_keys = {"formatter-cmds": ["black $file"]} for key in str_keys: diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 27f36ca67..b7861e779 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -58,8 +58,10 @@ class BestOptimization(BaseModel): candidate: OptimizedCandidate helper_functions: list[FunctionSource] runtime: int + replay_runtime: int | None winning_behavioral_test_results: TestResults winning_benchmarking_test_results: TestResults + winning_replay_benchmarking_test_results : TestResults | None = None class CodeString(BaseModel): @@ -198,7 +200,7 @@ class OriginalCodeBaseline(BaseModel): behavioral_test_results: TestResults benchmarking_test_results: TestResults runtime: int - coverage_results: Optional[CoverageData] + coverage_results: CoverageData | None class CoverageStatus(Enum): diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 66d3c6ab6..38277851b 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -117,7 +117,6 @@ def __init__( self.experiment_id = os.getenv("CODEFLASH_EXPERIMENT_ID", None) self.local_aiservice_client = LocalAiServiceClient() if self.experiment_id else None self.test_files = TestFiles(test_files=[]) - self.args = args # Check defaults for these self.function_trace_id: str = str(uuid.uuid4()) self.original_module_path = module_name_from_file_path(self.function_to_optimize.file_path, self.project_root) @@ -285,20 +284,22 @@ def optimize_function(self) -> Result[BestOptimization, str]: function_name=function_to_optimize_qualified_name, file_path=self.function_to_optimize.file_path, ) - speedup = explanation.speedup # eg. 1.2 means 1.2x faster + speedup = explanation.speedup # if self.args.benchmark: + original_replay_timing = original_code_baseline.benchmarking_test_results.total_replay_test_runtime() fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_file_name] for benchmark_name, og_benchmark_timing in fto_benchmark_timings.items(): print(f"Calculating speedup for benchmark {benchmark_name}") total_benchmark_timing = self.total_benchmark_timings[benchmark_name] # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values - expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + og_benchmark_timing / speedup + replay_speedup = original_replay_timing / best_optimization.replay_runtime - 1 + print(f"Replay speedup: {replay_speedup}") + expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / (replay_speedup + 1) * og_benchmark_timing print(f"Expected new benchmark timing: {expected_new_benchmark_timing}") print(f"Original benchmark timing: {total_benchmark_timing}") - print(f"Benchmark speedup: {total_benchmark_timing / expected_new_benchmark_timing}") - - speedup = total_benchmark_timing / expected_new_benchmark_timing - print(f"Speedup: {speedup}") + benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing + benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 + print(f"Benchmark speedup: {benchmark_speedup_percent:.2f}%") self.log_successful_optimization(explanation, generated_tests) @@ -447,13 +448,30 @@ def determine_best_candidate( ) tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") - + if self.args.benchmark: + original_code_replay_runtime = original_code_baseline.benchmarking_test_results.total_replay_test_runtime() + candidate_replay_runtime = candidate_result.benchmarking_test_results.total_replay_test_runtime() + replay_perf_gain = performance_gain( + original_runtime_ns=original_code_replay_runtime, + optimized_runtime_ns=candidate_replay_runtime, + ) + tree.add("Replay Benchmarking： ") + tree.add(f"Original summed runtime: {humanize_runtime(original_code_replay_runtime)}") + tree.add( + f"Best summed runtime: {humanize_runtime(candidate_replay_runtime)} " + f"(measured over {candidate_result.max_loop_count} " + f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + ) + tree.add(f"Speedup percentage: {replay_perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio: {replay_perf_gain + 1:.1f}X") best_optimization = BestOptimization( candidate=candidate, helper_functions=code_context.helper_functions, runtime=best_test_runtime, + replay_runtime=candidate_replay_runtime if self.args.benchmark else None, winning_behavioral_test_results=candidate_result.behavior_test_results, winning_benchmarking_test_results=candidate_result.benchmarking_test_results, + winning_replay_benchmarking_test_results=candidate_result.benchmarking_test_results, ) best_runtime_until_now = best_test_runtime else: @@ -664,6 +682,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi existing_test_files_count += 1 elif test_type == TestType.REPLAY_TEST: replay_test_files_count += 1 + print("Replay test found") elif test_type == TestType.CONCOLIC_COVERAGE_TEST: concolic_coverage_test_files_count += 1 else: @@ -708,6 +727,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi unique_instrumented_test_files.add(new_behavioral_test_path) unique_instrumented_test_files.add(new_perf_test_path) + if not self.test_files.get_by_original_file_path(path_obj_test_file): self.test_files.add( TestFile( @@ -719,6 +739,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi tests_in_file=[t.tests_in_file for t in tests_in_file_list], ) ) + logger.info( f"Discovered {existing_test_files_count} existing unit test file" f"{'s' if existing_test_files_count != 1 else ''}, {replay_test_files_count} replay test file" @@ -888,7 +909,6 @@ def establish_original_code_baseline( enable_coverage=False, code_context=code_context, ) - else: benchmarking_results = TestResults() start_time: float = time.time() @@ -917,7 +937,6 @@ def establish_original_code_baseline( ) console.rule() - total_timing = benchmarking_results.total_passed_runtime() # caution: doesn't handle the loop index functions_to_remove = [ result.id.test_function_name @@ -944,6 +963,9 @@ def establish_original_code_baseline( ) console.rule() logger.debug(f"Total original code runtime (ns): {total_timing}") + + if self.args.benchmark: + logger.info(f"Total replay test runtime: {humanize_runtime(benchmarking_results.total_replay_test_runtime())}") return Success( ( OriginalCodeBaseline( @@ -1062,6 +1084,15 @@ def run_optimized_candidate( console.rule() logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") + if self.args.benchmark: + total_candidate_replay_timing = ( + candidate_benchmarking_results.total_replay_test_runtime() + if candidate_benchmarking_results + else 0 + ) + logger.debug( + f"Total optimized code {optimization_candidate_index} replay benchmark runtime (ns): {total_candidate_replay_timing}" + ) return Success( OptimizedCandidateResult( max_loop_count=loop_count, diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 01a196143..9c5bc08ce 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -9,6 +9,7 @@ from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest +from codeflash.benchmarking.utils import print_benchmark_table from codeflash.cli_cmds.console import console, logger from codeflash.code_utils import env_utils from codeflash.code_utils.code_replacer import normalize_code, normalize_node @@ -23,7 +24,7 @@ from codeflash.verification.test_results import TestType 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 if TYPE_CHECKING: from argparse import Namespace @@ -98,11 +99,8 @@ def run(self) -> None: 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) - print(function_benchmark_timings) total_benchmark_timings = get_benchmark_timings(trace_dir) - print("Total benchmark timings:") - print(total_benchmark_timings) - # for function in fully_qualified_function_names: + print_benchmark_table(function_benchmark_timings, total_benchmark_timings) optimizations_found: int = 0 @@ -127,6 +125,7 @@ def run(self) -> None: console.rule() ph("cli-optimize-discovered-tests", {"num_tests": num_discovered_tests}) + for original_module_path in file_to_funcs_to_optimize: logger.info(f"Examining file {original_module_path!s}…") console.rule() @@ -217,6 +216,10 @@ def run(self) -> None: test_file.instrumented_behavior_file_path.unlink(missing_ok=True) if function_optimizer.test_cfg.concolic_test_root_dir: shutil.rmtree(function_optimizer.test_cfg.concolic_test_root_dir, ignore_errors=True) + if self.args.benchmark: + trace_dir = Path(self.args.benchmarks_root) / "codeflash_replay_tests" + if trace_dir.exists(): + shutil.rmtree(trace_dir, ignore_errors=True) if hasattr(get_run_tmp_file, "tmpdir"): get_run_tmp_file.tmpdir.cleanup() diff --git a/codeflash/tracer.py b/codeflash/tracer.py index 5bc1ae482..02a0e4157 100644 --- a/codeflash/tracer.py +++ b/codeflash/tracer.py @@ -58,6 +58,7 @@ def __init__( config_file_path: Path | None = None, max_function_count: int = 256, timeout: int | None = None, # seconds + benchmark: bool = False, ) -> None: """:param output: The path to the output trace file :param functions: List of functions to trace. If None, trace all functions @@ -95,7 +96,6 @@ def __init__( self.max_function_count = max_function_count self.config, found_config_path = parse_config_file(config_file_path) self.project_root = project_root_from_module_root(Path(self.config["module_root"]), found_config_path) - print("project_root", self.project_root) self.ignored_functions = {"", "", "", "", "", ""} self.file_being_called_from: str = str(Path(sys._getframe().f_back.f_code.co_filename).name).replace(".", "_") @@ -105,6 +105,7 @@ def __init__( self.next_insert = 1000 self.trace_count = 0 + self.benchmark = benchmark # Profiler variables self.bias = 0 # calibration constant self.timings = {} @@ -184,18 +185,25 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: cur.execute("INSERT INTO total_time VALUES (?)", (self.total_tt,)) self.con.commit() self.con.close() + function_string = [str(function.file_name) + ":" + (function.class_name + ":" if function.class_name else "") + function.function_name for function in self.function_modules] + # print(function_string) # filter any functions where we did not capture the return + # self.function_modules = [ + # function + # for function in self.function_modules + # if self.function_count[ + # str(function.file_name) + # + ":" + # + (function.class_name + ":" if function.class_name else "") + # + function.function_name + # ] + # > 0 + # ] self.function_modules = [ function for function in self.function_modules - if self.function_count[ - str(function.file_name) - + ":" - + (function.class_name + ":" if function.class_name else "") - + function.function_name - ] - > 0 + if str(str(function.file_name) + ":" + (function.class_name + ":" if function.class_name else "") + function.function_name) in self.function_count ] replay_test = create_trace_replay_test( @@ -207,13 +215,21 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: # Need a better way to store the replay test # function_path = "_".join(self.functions) if self.functions else self.file_being_called_from function_path = self.file_being_called_from - test_file_path = get_test_file_path( - test_dir=Path(self.config["tests_root"]), function_name=function_path, test_type="replay" - ) + if self.benchmark and self.config["benchmarks_root"]: + # check if replay test dir exists, create + replay_test_dir = Path(self.config["benchmarks_root"]) / "codeflash_replay_tests" + if not replay_test_dir.exists(): + replay_test_dir.mkdir(parents=True) + test_file_path = get_test_file_path( + test_dir=replay_test_dir, function_name=function_path, test_type="replay" + ) + else: + test_file_path = get_test_file_path( + test_dir=Path(self.config["tests_root"]), function_name=function_path, test_type="replay" + ) replay_test = isort.code(replay_test) with open(test_file_path, "w", encoding="utf8") as file: file.write(replay_test) - console.print( f"Codeflash: Traced {self.trace_count} function calls successfully and replay test created at - {test_file_path}", crop=False, diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py index db01ff049..916f6da11 100644 --- a/codeflash/verification/test_results.py +++ b/codeflash/verification/test_results.py @@ -193,6 +193,35 @@ def total_passed_runtime(self) -> int: ] ) + def usable_replay_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: + """Collect runtime data for replay tests that passed and have runtime information. + + :return: A dictionary mapping invocation IDs to lists of runtime values. + """ + usable_runtimes = [ + (result.id, result.runtime) + for result in self.test_results + if result.did_pass and result.runtime and result.test_type == TestType.REPLAY_TEST + ] + + return { + usable_id: [runtime[1] for runtime in usable_runtimes if runtime[0] == usable_id] + for usable_id in {runtime[0] for runtime in usable_runtimes} + } + + def total_replay_test_runtime(self) -> int: + """Calculate the sum of runtimes of replay test cases that passed, where a testcase runtime + is the minimum value of all looped execution runtimes. + + :return: The runtime in nanoseconds. + """ + replay_runtime_data = self.usable_replay_runtime_data_by_test_case() + + return sum([ + min(runtimes) + for invocation_id, runtimes in replay_runtime_data.items() + ]) if replay_runtime_data else 0 + def __iter__(self) -> Iterator[FunctionTestInvocation]: return iter(self.test_results) diff --git a/codeflash/verification/test_runner.py b/codeflash/verification/test_runner.py index 46203e65a..e85763c08 100644 --- a/codeflash/verification/test_runner.py +++ b/codeflash/verification/test_runner.py @@ -56,6 +56,8 @@ def run_behavioral_tests( "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", + "-o", + "addopts=", "--codeflash_loops_scope=session", "--codeflash_min_loops=1", "--codeflash_max_loops=1", From 6b4b68a93d055607cf03566c0eb01833a6d0163e Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 11 Mar 2025 14:37:44 -0700 Subject: [PATCH 004/122] basic version working on bubble sort --- code_to_optimize/bubble_sort.py | 12 ++----- .../benchmarks/test_process_and_sort.py | 2 +- .../tests/unittest/test_bubble_sort.py | 36 +++++++++---------- .../unittest/test_bubble_sort_parametrized.py | 36 +++++++++---------- codeflash/optimization/optimizer.py | 2 ++ pyproject.toml | 5 +-- 6 files changed, 45 insertions(+), 48 deletions(-) diff --git a/code_to_optimize/bubble_sort.py b/code_to_optimize/bubble_sort.py index 787cc4a90..fd53c04a7 100644 --- a/code_to_optimize/bubble_sort.py +++ b/code_to_optimize/bubble_sort.py @@ -1,10 +1,4 @@ def sorter(arr): - print("codeflash stdout: Sorting list") - for i in range(len(arr)): - for j in range(len(arr) - 1): - if arr[j] > arr[j + 1]: - temp = arr[j] - arr[j] = arr[j + 1] - arr[j + 1] = temp - print(f"result: {arr}") - return arr \ No newline at end of file + # Utilizing Python's built-in Timsort algorithm for better performance + arr.sort() + return arr diff --git a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py index ca2f0ef65..93d78afef 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py @@ -1,5 +1,5 @@ from code_to_optimize.process_and_bubble_sort import compute_and_sort -from code_to_optimize.bubble_sort2 import sorter +from code_to_optimize.bubble_sort import sorter def test_compute_and_sort(benchmark): result = benchmark(compute_and_sort, list(reversed(range(5000)))) assert result == 6247083.5 diff --git a/code_to_optimize/tests/unittest/test_bubble_sort.py b/code_to_optimize/tests/unittest/test_bubble_sort.py index 200f82b7a..4c76414ef 100644 --- a/code_to_optimize/tests/unittest/test_bubble_sort.py +++ b/code_to_optimize/tests/unittest/test_bubble_sort.py @@ -1,18 +1,18 @@ -import unittest - -from code_to_optimize.bubble_sort import sorter - - -class TestPigLatin(unittest.TestCase): - def test_sort(self): - input = [5, 4, 3, 2, 1, 0] - output = sorter(input) - self.assertEqual(output, [0, 1, 2, 3, 4, 5]) - - input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0] - output = sorter(input) - self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) - - input = list(reversed(range(5000))) - output = sorter(input) - self.assertEqual(output, list(range(5000))) +# import unittest +# +# from code_to_optimize.bubble_sort import sorter +# +# +# class TestPigLatin(unittest.TestCase): +# def test_sort(self): +# input = [5, 4, 3, 2, 1, 0] +# output = sorter(input) +# self.assertEqual(output, [0, 1, 2, 3, 4, 5]) +# +# input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0] +# output = sorter(input) +# self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) +# +# input = list(reversed(range(5000))) +# output = sorter(input) +# self.assertEqual(output, list(range(5000))) diff --git a/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py b/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py index 59c86abc8..c1aef993b 100644 --- a/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py +++ b/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py @@ -1,18 +1,18 @@ -import unittest - -from parameterized import parameterized - -from code_to_optimize.bubble_sort import sorter - - -class TestPigLatin(unittest.TestCase): - @parameterized.expand( - [ - ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]), - ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), - (list(reversed(range(50))), list(range(50))), - ] - ) - def test_sort(self, input, expected_output): - output = sorter(input) - self.assertEqual(output, expected_output) +# import unittest +# +# from parameterized import parameterized +# +# from code_to_optimize.bubble_sort import sorter +# +# +# class TestPigLatin(unittest.TestCase): +# @parameterized.expand( +# [ +# ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]), +# ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), +# (list(reversed(range(50))), list(range(50))), +# ] +# ) +# def test_sort(self, input, expected_output): +# output = sorter(input) +# self.assertEqual(output, expected_output) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index de2f6deed..5a1416316 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -78,6 +78,8 @@ def run(self) -> None: function_optimizer = None file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]] num_optimizable_functions: int + # if self.args.benchmark: + # discover functions (file_to_funcs_to_optimize, num_optimizable_functions) = get_functions_to_optimize( optimize_all=self.args.all, replay_test=self.args.replay_test, diff --git a/pyproject.toml b/pyproject.toml index 2e71f2a0a..877815004 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -216,8 +216,9 @@ initial-content = """ [tool.codeflash] -module-root = "codeflash" -tests-root = "tests" +module-root = "code_to_optimize" +tests-root = "code_to_optimize/tests" +benchmarks-root = "code_to_optimize/tests/pytest/benchmarks" test-framework = "pytest" formatter-cmds = [ "uvx ruff check --exit-zero --fix $file", From 887e3cba56b029ce281f95840f2e967b9bef5099 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 11 Mar 2025 16:00:21 -0700 Subject: [PATCH 005/122] initial attempt for codeflash_trace_decorator --- .../benchmarking/codeflash_trace_decorator.py | 73 +++++++++++++++++++ .../instrument_codeflash_trace.py | 1 + codeflash/benchmarking/plugin/plugin.py | 46 ++---------- .../pytest_new_process_trace_benchmarks.py | 3 +- codeflash/benchmarking/trace_benchmarks.py | 5 +- codeflash/optimization/optimizer.py | 35 ++++++--- tests/test_trace_benchmarks.py | 2 +- 7 files changed, 106 insertions(+), 59 deletions(-) create mode 100644 codeflash/benchmarking/codeflash_trace_decorator.py create mode 100644 codeflash/benchmarking/instrument_codeflash_trace.py diff --git a/codeflash/benchmarking/codeflash_trace_decorator.py b/codeflash/benchmarking/codeflash_trace_decorator.py new file mode 100644 index 000000000..f996fa295 --- /dev/null +++ b/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 diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -0,0 +1 @@ + diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index 80accec22..bb903e554 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/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() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 6b91e2b4f..85a6755bf 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/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}") diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 2d3acdd66..bec5a03d4 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/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, diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 5a1416316..1d03de150 100644 --- a/codeflash/optimization/optimizer.py +++ b/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 diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index acd40b0b3..688e23d8c 100644 --- a/tests/test_trace_benchmarks.py +++ b/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"]) \ No newline at end of file + trace_benchmarks_pytest(benchmarks_root, project_root) \ No newline at end of file From 84bd0f09dee4c5d20c499550e375af1145608a63 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 11 Mar 2025 17:35:34 -0700 Subject: [PATCH 006/122] improvements --- code_to_optimize/bubble_sort.py | 10 +- codeflash/benchmarking/codeflash_trace.py | 122 ++++++++++++++++++ .../benchmarking/codeflash_trace_decorator.py | 73 ----------- codeflash/benchmarking/plugin/plugin.py | 8 +- tests/test_codeflash_trace_decorator.py | 15 +++ 5 files changed, 150 insertions(+), 78 deletions(-) create mode 100644 codeflash/benchmarking/codeflash_trace.py delete mode 100644 codeflash/benchmarking/codeflash_trace_decorator.py create mode 100644 tests/test_codeflash_trace_decorator.py diff --git a/code_to_optimize/bubble_sort.py b/code_to_optimize/bubble_sort.py index fd53c04a7..41cca9cea 100644 --- a/code_to_optimize/bubble_sort.py +++ b/code_to_optimize/bubble_sort.py @@ -1,4 +1,10 @@ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +@codeflash_trace("bubble_sort.trace") def sorter(arr): - # Utilizing Python's built-in Timsort algorithm for better performance - arr.sort() + for i in range(len(arr)): + for j in range(len(arr) - 1): + if arr[j] > arr[j + 1]: + temp = arr[j] + arr[j] = arr[j + 1] + arr[j + 1] = temp return arr diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py new file mode 100644 index 000000000..428f4a22c --- /dev/null +++ b/codeflash/benchmarking/codeflash_trace.py @@ -0,0 +1,122 @@ +import functools +import os +import pickle +import sqlite3 +import time +from typing import Callable + + +class CodeflashTrace: + """A class that provides both a decorator for tracing function calls + and a context manager for managing the tracing data lifecycle. + """ + + def __init__(self) -> None: + self.function_calls_data = [] + + def __enter__(self) -> None: + # Initialize for context manager use + self.function_calls_data = [] + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + # Cleanup is optional here + pass + + def __call__(self, func: Callable) -> Callable: + """Use as a decorator to trace function execution. + + Args: + func: The function to be decorated + + Returns: + The wrapped function + + """ + @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() + overhead_time = 0 + + try: + # 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", "") + + # Calculate overhead time + overhead_end_time = time.time() + overhead_time = overhead_end_time - overhead_start_time + + self.function_calls_data.append( + (func.__name__, func.__module__, func.__code__.co_filename, + benchmark_function_name, benchmark_file_name, execution_time, + overhead_time, pickled_args, pickled_kwargs) + ) + + except Exception as e: + print(f"Error in codeflash_trace: {e}") + + return result + return wrapper + + def write_to_db(self, output_file: str) -> None: + """Write all collected function call data to the SQLite database. + + Args: + output_file: Path to the SQLite database file where results will be stored + + """ + if not self.function_calls_data: + print("No function call data to write") + return + + 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, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" + ) + + # Insert all data at once + cur.executemany( + "INSERT INTO function_calls " + "(function_name, class_name, file_name, benchmark_function_name, " + "benchmark_file_name, time_ns, overhead_time_ns, args, kwargs) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + self.function_calls_data + ) + + # Commit and close + con.commit() + con.close() + + print(f"Successfully wrote {len(self.function_calls_data)} function call records to {output_file}") + + # Clear the data after writing + self.function_calls_data.clear() + + except Exception as e: + print(f"Error writing function calls to database: {e}") + +# Create a singleton instance +codeflash_trace = CodeflashTrace() diff --git a/codeflash/benchmarking/codeflash_trace_decorator.py b/codeflash/benchmarking/codeflash_trace_decorator.py deleted file mode 100644 index f996fa295..000000000 --- a/codeflash/benchmarking/codeflash_trace_decorator.py +++ /dev/null @@ -1,73 +0,0 @@ -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 diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index bb903e554..6d8db9bf9 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -1,6 +1,6 @@ import pytest import time - +import os class CodeFlashPlugin: @staticmethod def pytest_addoption(parser): @@ -35,9 +35,11 @@ def benchmark(request): class Benchmark: def __call__(self, func, *args, **kwargs): - start = time.time_ns() + os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = request.node.name + os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = request.node.fspath.basename + start = time.process_time_ns() result = func(*args, **kwargs) - end = time.time_ns() + end = time.process_time_ns() print(f"Benchmark: {func.__name__} took {end - start} ns") return result diff --git a/tests/test_codeflash_trace_decorator.py b/tests/test_codeflash_trace_decorator.py new file mode 100644 index 000000000..251b668ec --- /dev/null +++ b/tests/test_codeflash_trace_decorator.py @@ -0,0 +1,15 @@ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +from pathlib import Path + +@codeflash_trace("test_codeflash_trace.trace") +def example_function(arr): + arr.sort() + return arr + + +def test_codeflash_trace_decorator(): + arr = [3, 1, 2] + result = example_function(arr) + # cleanup test trace file using Path + assert result == [1, 2, 3] + Path("test_codeflash_trace.trace").unlink() \ No newline at end of file From c4694b74dfe042e01046fcd044a9f24394ab9bdd Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 12 Mar 2025 11:46:29 -0700 Subject: [PATCH 007/122] work on new replay_test logic --- code_to_optimize/bubble_sort.py | 2 +- codeflash/benchmarking/codeflash_trace.py | 70 +++++++- .../pytest_new_process_trace_benchmarks.py | 5 + codeflash/benchmarking/replay_test.py | 159 ++++++++++++++++++ codeflash/benchmarking/trace_benchmarks.py | 3 +- codeflash/optimization/optimizer.py | 7 + tests/test_codeflash_trace_decorator.py | 4 +- tests/test_trace_benchmarks.py | 6 +- 8 files changed, 246 insertions(+), 10 deletions(-) create mode 100644 codeflash/benchmarking/replay_test.py diff --git a/code_to_optimize/bubble_sort.py b/code_to_optimize/bubble_sort.py index 41cca9cea..91b77f50c 100644 --- a/code_to_optimize/bubble_sort.py +++ b/code_to_optimize/bubble_sort.py @@ -1,5 +1,5 @@ from codeflash.benchmarking.codeflash_trace import codeflash_trace -@codeflash_trace("bubble_sort.trace") +@codeflash_trace def sorter(arr): for i in range(len(arr)): for j in range(len(arr) - 1): diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 428f4a22c..45c7fa6c2 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -14,10 +14,10 @@ class CodeflashTrace: def __init__(self) -> None: self.function_calls_data = [] - def __enter__(self) -> None: - # Initialize for context manager use - self.function_calls_data = [] - return self + # def __enter__(self) -> None: + # # Initialize for context manager use + # self.function_calls_data = [] + # return self def __exit__(self, exc_type, exc_val, exc_tb) -> None: # Cleanup is optional here @@ -82,7 +82,7 @@ def write_to_db(self, output_file: str) -> None: if not self.function_calls_data: print("No function call data to write") return - + self.db_path = output_file try: # Connect to the database con = sqlite3.connect(output_file) @@ -118,5 +118,65 @@ def write_to_db(self, output_file: str) -> None: except Exception as e: print(f"Error writing function calls to database: {e}") + def print_codeflash_db(self, limit: int = None) -> None: + """ + Print the contents of a CodeflashTrace SQLite database. + + Args: + db_path: Path to the SQLite database file + limit: Maximum number of records to print (None for all) + """ + try: + # Connect to the database + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + + # Get the count of records + cursor.execute("SELECT COUNT(*) FROM function_calls") + total_records = cursor.fetchone()[0] + print(f"Found {total_records} function call records in {self.db_path}") + + # Build the query with optional limit + query = "SELECT * FROM function_calls" + if limit: + query += f" LIMIT {limit}" + + # Execute the query + cursor.execute(query) + + # Print column names + columns = [desc[0] for desc in cursor.description] + print("\nColumns:", columns) + print("\n" + "=" * 80 + "\n") + + # Print each row + for i, row in enumerate(cursor.fetchall()): + print(f"Record #{i + 1}:") + print(f" Function: {row[0]}") + print(f" Module: {row[1]}") + print(f" File: {row[2]}") + print(f" Benchmark Function: {row[3] or 'N/A'}") + print(f" Benchmark File: {row[4] or 'N/A'}") + print(f" Execution Time: {row[5]:.6f} seconds") + print(f" Overhead Time: {row[6]:.6f} seconds") + + # Unpickle and print args and kwargs + try: + args = pickle.loads(row[7]) + kwargs = pickle.loads(row[8]) + + print(f" Args: {args}") + print(f" Kwargs: {kwargs}") + except Exception as e: + print(f" Error unpickling args/kwargs: {e}") + print(f" Raw args: {row[7]}") + print(f" Raw kwargs: {row[8]}") + + print("\n" + "-" * 40 + "\n") + + conn.close() + + except Exception as e: + print(f"Error reading database: {e}") # Create a singleton instance codeflash_trace = CodeflashTrace() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 85a6755bf..ebe1fa4ae 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -1,7 +1,10 @@ import sys from plugin.plugin import CodeFlashPlugin +from codeflash.benchmarking.codeflash_trace import codeflash_trace +from codeflash.code_utils.code_utils import get_run_tmp_file benchmarks_root = sys.argv[1] +output_file = sys.argv[2] if __name__ == "__main__": import pytest @@ -9,6 +12,8 @@ exitcode = pytest.main( [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashPlugin()] ) + codeflash_trace.write_to_db(output_file) + codeflash_trace.print_codeflash_db() except Exception as e: print(f"Failed to collect tests: {e!s}") exitcode = -1 \ No newline at end of file diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py new file mode 100644 index 000000000..0bc2de1d4 --- /dev/null +++ b/codeflash/benchmarking/replay_test.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +import sqlite3 +import textwrap +from collections.abc import Generator +from typing import Any, List, Optional + +from codeflash.discovery.functions_to_optimize import FunctionProperties, inspect_top_level_functions_or_methods +from codeflash.tracing.tracing_utils import FunctionModules + + +def get_next_arg_and_return( + trace_file: str, function_name: str, file_name: str, class_name: str | None, num_to_get: int = 25 +) -> Generator[Any]: + db = sqlite3.connect(trace_file) + cur = db.cursor() + limit = num_to_get + if class_name is not None: + cursor = cur.execute( + "SELECT * FROM function_calls WHERE function_name = ? AND file_namename = ? AND class_name = ? ORDER BY time_ns ASC LIMIT ?", + (function_name, file_name, class_name, limit), + ) + else: + cursor = cur.execute( + "SELECT * FROM function_calls WHERE function_name = ? AND file_namename = ? ORDER BY time_ns ASC LIMIT ?", + (function_name, file_name, limit), + ) + + while (val := cursor.fetchone()) is not None: + yield val[8], val[9] + + +def get_function_alias(module: str, function_name: str) -> str: + return "_".join(module.split(".")) + "_" + function_name + + +def create_trace_replay_test( + trace_file: str, functions: list[FunctionModules], test_framework: str = "pytest", max_run_count=100 +) -> str: + assert test_framework in ["pytest", "unittest"] + + imports = f"""import dill as pickle +{"import unittest" if test_framework == "unittest" else ""} +from codeflash.tracing.replay_test import get_next_arg_and_return +""" + + # TODO: Module can have "-" character if the module-root is ".". Need to handle that case + function_properties: list[FunctionProperties] = [ + inspect_top_level_functions_or_methods( + file_name=function.file_name, + function_or_method_name=function.function_name, + class_name=function.class_name, + line_no=function.line_no, + ) + for function in functions + ] + function_imports = [] + for function, function_property in zip(functions, function_properties): + if not function_property.is_top_level: + # can't be imported and run in the replay test + continue + if function_property.is_staticmethod: + function_imports.append( + f"from {function.module_name} import {function_property.staticmethod_class_name} as {get_function_alias(function.module_name, function_property.staticmethod_class_name)}" + ) + elif function.class_name: + function_imports.append( + f"from {function.module_name} import {function.class_name} as {get_function_alias(function.module_name, function.class_name)}" + ) + else: + function_imports.append( + f"from {function.module_name} import {function.function_name} as {get_function_alias(function.module_name, function.function_name)}" + ) + + imports += "\n".join(function_imports) + functions_to_optimize = [function.function_name for function in functions if function.function_name != "__init__"] + metadata = f"""functions = {functions_to_optimize} +trace_file_path = r"{trace_file}" +""" # trace_file_path path is parsed with regex later, format is important + test_function_body = textwrap.dedent( + """\ + for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): + args = pickle.loads(arg_val_pkl) + ret = {function_name}({args}) + """ + ) + test_class_method_body = textwrap.dedent( + """\ + for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + args = pickle.loads(arg_val_pkl){filter_variables} + ret = {class_name_alias}{method_name}(**args) + """ + ) + test_class_staticmethod_body = textwrap.dedent( + """\ + for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): + args = pickle.loads(arg_val_pkl){filter_variables} + ret = {class_name_alias}{method_name}(**args) + """ + ) + if test_framework == "unittest": + self = "self" + test_template = "\nclass TestTracedFunctions(unittest.TestCase):\n" + else: + test_template = "" + self = "" + for func, func_property in zip(functions, function_properties): + if not func_property.is_top_level: + # can't be imported and run in the replay test + continue + if func.class_name is None and not func_property.is_staticmethod: + alias = get_function_alias(func.module_name, func.function_name) + test_body = test_function_body.format( + function_name=alias, + file_name=func.file_name, + orig_function_name=func.function_name, + max_run_count=max_run_count, + args="**args" if func_property.has_args else "", + ) + elif func_property.is_staticmethod: + class_name_alias = get_function_alias(func.module_name, func_property.staticmethod_class_name) + alias = get_function_alias( + func.module_name, func_property.staticmethod_class_name + "_" + func.function_name + ) + method_name = "." + func.function_name if func.function_name != "__init__" else "" + test_body = test_class_staticmethod_body.format( + orig_function_name=func.function_name, + file_name=func.file_name, + class_name_alias=class_name_alias, + method_name=method_name, + max_run_count=max_run_count, + filter_variables="", + ) + else: + class_name_alias = get_function_alias(func.module_name, func.class_name) + alias = get_function_alias(func.module_name, func.class_name + "_" + func.function_name) + + if func_property.is_classmethod: + filter_variables = '\n args.pop("cls", None)' + elif func.function_name == "__init__": + filter_variables = '\n args.pop("__class__", None)' + else: + filter_variables = "" + method_name = "." + func.function_name if func.function_name != "__init__" else "" + test_body = test_class_method_body.format( + orig_function_name=func.function_name, + file_name=func.file_name, + class_name_alias=class_name_alias, + class_name=func.class_name, + method_name=method_name, + max_run_count=max_run_count, + filter_variables=filter_variables, + ) + formatted_test_body = textwrap.indent(test_body, " " if test_framework == "unittest" else " ") + + test_template += " " if test_framework == "unittest" else "" + test_template += f"def test_{alias}({self}):\n{formatted_test_body}\n" + + return imports + "\n" + metadata + "\n" + test_template diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index bec5a03d4..54e0b5118 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -3,12 +3,13 @@ from pathlib import Path import subprocess -def trace_benchmarks_pytest(benchmarks_root: Path, project_root: Path) -> None: +def trace_benchmarks_pytest(benchmarks_root: Path, project_root: Path, output_file: Path) -> None: result = subprocess.run( [ SAFE_SYS_EXECUTABLE, Path(__file__).parent / "pytest_new_process_trace_benchmarks.py", benchmarks_root, + output_file, ], cwd=project_root, check=False, diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 1d03de150..26fc70aa7 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -25,7 +25,10 @@ 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 codeflash.benchmarking.codeflash_trace import codeflash_trace + from collections import defaultdict + if TYPE_CHECKING: from argparse import Namespace @@ -104,12 +107,16 @@ def run(self) -> None: #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") + except Exception as e: + logger.info(f"Error while tracing existing benchmarks: {e}") + logger.info(f"Information on existing benchmarks will not be available for this run.") 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]) + codeflash_trace.print_trace_info() # 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) diff --git a/tests/test_codeflash_trace_decorator.py b/tests/test_codeflash_trace_decorator.py index 251b668ec..37234d85a 100644 --- a/tests/test_codeflash_trace_decorator.py +++ b/tests/test_codeflash_trace_decorator.py @@ -1,7 +1,8 @@ from codeflash.benchmarking.codeflash_trace import codeflash_trace from pathlib import Path +from codeflash.code_utils.code_utils import get_run_tmp_file -@codeflash_trace("test_codeflash_trace.trace") +@codeflash_trace def example_function(arr): arr.sort() return arr @@ -12,4 +13,3 @@ def test_codeflash_trace_decorator(): result = example_function(arr) # cleanup test trace file using Path assert result == [1, 2, 3] - Path("test_codeflash_trace.trace").unlink() \ No newline at end of file diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 688e23d8c..071535c6a 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,8 +1,12 @@ from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from pathlib import Path +from codeflash.code_utils.code_utils import get_run_tmp_file 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) \ No newline at end of file + output_file = Path("test_trace_benchmarks.trace").resolve() + trace_benchmarks_pytest(benchmarks_root, project_root, output_file) + assert output_file.exists() + output_file.unlink() \ No newline at end of file From 1801d414bfab068757b6b736393e1f042cfce664 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 13 Mar 2025 18:14:38 -0700 Subject: [PATCH 008/122] initial replay test version working --- code_to_optimize/bubble_sort.py | 2 - .../benchmarks/test_benchmark_bubble_sort.py | 2 +- .../benchmarks/test_process_and_sort.py | 4 +- codeflash/benchmarking/codeflash_trace.py | 122 ++++- .../pytest_new_process_trace_benchmarks.py | 9 +- codeflash/benchmarking/replay_test.py | 149 +++-- codeflash/benchmarking/trace_benchmarks.py | 3 +- codeflash/optimization/function_optimizer.py | 4 +- codeflash/tracer.py | 511 +++++++++++------- tests/test_trace_benchmarks.py | 21 +- 10 files changed, 525 insertions(+), 302 deletions(-) diff --git a/code_to_optimize/bubble_sort.py b/code_to_optimize/bubble_sort.py index 91b77f50c..db7db5f92 100644 --- a/code_to_optimize/bubble_sort.py +++ b/code_to_optimize/bubble_sort.py @@ -1,5 +1,3 @@ -from codeflash.benchmarking.codeflash_trace import codeflash_trace -@codeflash_trace def sorter(arr): for i in range(len(arr)): for j in range(len(arr) - 1): diff --git a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py index f1ebcf5c7..21c2bbb29 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py @@ -1,6 +1,6 @@ import pytest -from code_to_optimize.bubble_sort import sorter +from code_to_optimize.bubble_sort_codeflash_trace import sorter def test_sort(benchmark): diff --git a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py index 93d78afef..2713721e4 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py @@ -1,5 +1,5 @@ -from code_to_optimize.process_and_bubble_sort import compute_and_sort -from code_to_optimize.bubble_sort import sorter +from code_to_optimize.process_and_bubble_sort_codeflash_trace import compute_and_sort +from code_to_optimize.bubble_sort_codeflash_trace import sorter def test_compute_and_sort(benchmark): result = benchmark(compute_and_sort, list(reversed(range(5000)))) assert result == 6247083.5 diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 45c7fa6c2..c678b7643 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -3,9 +3,12 @@ import pickle import sqlite3 import time +from pathlib import Path from typing import Callable + + class CodeflashTrace: """A class that provides both a decorator for tracing function calls and a context manager for managing the tracing data lifecycle. @@ -60,8 +63,12 @@ def wrapper(*args, **kwargs): overhead_end_time = time.time() overhead_time = overhead_end_time - overhead_start_time + class_name = "" + qualname = func.__qualname__ + if "." in qualname: + class_name = qualname.split(".")[0] self.function_calls_data.append( - (func.__name__, func.__module__, func.__code__.co_filename, + (func.__name__, class_name, func.__module__, func.__code__.co_filename, benchmark_function_name, benchmark_file_name, execution_time, overhead_time, pickled_args, pickled_kwargs) ) @@ -92,7 +99,7 @@ def write_to_db(self, output_file: str) -> None: # 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, " + "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," "benchmark_function_name TEXT, benchmark_file_name TEXT, " "time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" ) @@ -100,9 +107,9 @@ def write_to_db(self, output_file: str) -> None: # Insert all data at once cur.executemany( "INSERT INTO function_calls " - "(function_name, class_name, file_name, benchmark_function_name, " + "(function_name, class_name, module_name, file_name, benchmark_function_name, " "benchmark_file_name, time_ns, overhead_time_ns, args, kwargs) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", self.function_calls_data ) @@ -153,24 +160,25 @@ def print_codeflash_db(self, limit: int = None) -> None: for i, row in enumerate(cursor.fetchall()): print(f"Record #{i + 1}:") print(f" Function: {row[0]}") - print(f" Module: {row[1]}") - print(f" File: {row[2]}") - print(f" Benchmark Function: {row[3] or 'N/A'}") - print(f" Benchmark File: {row[4] or 'N/A'}") - print(f" Execution Time: {row[5]:.6f} seconds") - print(f" Overhead Time: {row[6]:.6f} seconds") + print(f" Class: {row[1]}") + print(f" Module: {row[2]}") + print(f" File: {row[3]}") + print(f" Benchmark Function: {row[4] or 'N/A'}") + print(f" Benchmark File: {row[5] or 'N/A'}") + print(f" Execution Time: {row[6]:.6f} seconds") + print(f" Overhead Time: {row[7]:.6f} seconds") # Unpickle and print args and kwargs try: - args = pickle.loads(row[7]) - kwargs = pickle.loads(row[8]) + args = pickle.loads(row[8]) + kwargs = pickle.loads(row[9]) print(f" Args: {args}") print(f" Kwargs: {kwargs}") except Exception as e: print(f" Error unpickling args/kwargs: {e}") - print(f" Raw args: {row[7]}") - print(f" Raw kwargs: {row[8]}") + print(f" Raw args: {row[8]}") + print(f" Raw kwargs: {row[9]}") print("\n" + "-" * 40 + "\n") @@ -178,5 +186,91 @@ def print_codeflash_db(self, limit: int = None) -> None: except Exception as e: print(f"Error reading database: {e}") + + def generate_replay_test(self, output_dir: str = None, project_root: str = "", test_framework: str = "pytest", + max_run_count: int = 100) -> None: + """ + Generate multiple replay tests from the traced function calls, grouping by benchmark name. + + Args: + output_dir: Directory to write the generated tests (if None, only returns the code) + project_root: Root directory of the project for module imports + test_framework: 'pytest' or 'unittest' + max_run_count: Maximum number of runs to include per function + + Returns: + Dictionary mapping benchmark names to generated test code + """ + import isort + from codeflash.verification.verification_utils import get_test_file_path + + if not self.db_path: + print("No database path set. Call write_to_db first or set db_path manually.") + return {} + + try: + # Import the function here to avoid circular imports + from codeflash.benchmarking.replay_test import create_trace_replay_test + + print("connecting to: ", self.db_path) + # Connect to the database + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + + # Get distinct benchmark names + cursor.execute( + "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM function_calls" + ) + benchmarks = cursor.fetchall() + + # Generate a test for each benchmark + for benchmark in benchmarks: + benchmark_function_name, benchmark_file_name = benchmark + # Get functions associated with this benchmark + cursor.execute( + "SELECT DISTINCT function_name, class_name, module_name, file_name FROM function_calls " + "WHERE benchmark_function_name = ? AND benchmark_file_name = ?", + (benchmark_function_name, benchmark_file_name) + ) + + functions_data = [] + for func_row in cursor.fetchall(): + function_name, class_name, module_name, file_name = func_row + + # Add this function to our list + functions_data.append({ + "function_name": function_name, + "class_name": class_name, + "file_name": file_name, + "module_name": module_name + }) + + if not functions_data: + print(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") + continue + + # Generate the test code for this benchmark + test_code = create_trace_replay_test( + trace_file=self.db_path, + functions_data=functions_data, + test_framework=test_framework, + max_run_count=max_run_count, + ) + test_code = isort.code(test_code) + + # Write to file if requested + if output_dir: + output_file = get_test_file_path( + test_dir=Path(output_dir), function_name=f"{benchmark_file_name[5:]}_{benchmark_function_name}", test_type="replay" + ) + with open(output_file, 'w') as f: + f.write(test_code) + print(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") + + conn.close() + + except Exception as e: + print(f"Error generating replay tests: {e}") + # Create a singleton instance codeflash_trace = CodeflashTrace() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index ebe1fa4ae..8e1958fec 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -1,10 +1,16 @@ import sys +from pathlib import Path + +from codeflash.verification.verification_utils import get_test_file_path from plugin.plugin import CodeFlashPlugin from codeflash.benchmarking.codeflash_trace import codeflash_trace from codeflash.code_utils.code_utils import get_run_tmp_file benchmarks_root = sys.argv[1] -output_file = sys.argv[2] +tests_root = sys.argv[2] +output_file = sys.argv[3] +# current working directory +project_root = Path.cwd() if __name__ == "__main__": import pytest @@ -14,6 +20,7 @@ ) codeflash_trace.write_to_db(output_file) codeflash_trace.print_codeflash_db() + codeflash_trace.generate_replay_test(tests_root, project_root, test_framework="pytest") except Exception as e: print(f"Failed to collect tests: {e!s}") exitcode = -1 \ No newline at end of file diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 0bc2de1d4..9bc2c79f3 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -3,31 +3,29 @@ import sqlite3 import textwrap from collections.abc import Generator -from typing import Any, List, Optional - -from codeflash.discovery.functions_to_optimize import FunctionProperties, inspect_top_level_functions_or_methods -from codeflash.tracing.tracing_utils import FunctionModules +from typing import Any, Dict def get_next_arg_and_return( - trace_file: str, function_name: str, file_name: str, class_name: str | None, num_to_get: int = 25 + trace_file: str, function_name: str, file_name: str, class_name: str | None = None, num_to_get: int = 25 ) -> Generator[Any]: db = sqlite3.connect(trace_file) cur = db.cursor() limit = num_to_get + if class_name is not None: cursor = cur.execute( - "SELECT * FROM function_calls WHERE function_name = ? AND file_namename = ? AND class_name = ? ORDER BY time_ns ASC LIMIT ?", + "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? AND class_name = ? ORDER BY time_ns ASC LIMIT ?", (function_name, file_name, class_name, limit), ) else: cursor = cur.execute( - "SELECT * FROM function_calls WHERE function_name = ? AND file_namename = ? ORDER BY time_ns ASC LIMIT ?", + "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? ORDER BY time_ns ASC LIMIT ?", (function_name, file_name, limit), ) while (val := cursor.fetchone()) is not None: - yield val[8], val[9] + yield val[8], val[9] # args and kwargs are at indices 7 and 8 def get_function_alias(module: str, function_name: str) -> str: @@ -35,122 +33,109 @@ def get_function_alias(module: str, function_name: str) -> str: def create_trace_replay_test( - trace_file: str, functions: list[FunctionModules], test_framework: str = "pytest", max_run_count=100 + trace_file: str, + functions_data: list[Dict[str, Any]], + test_framework: str = "pytest", + max_run_count=100 ) -> str: + """Create a replay test for functions based on trace data. + + Args: + trace_file: Path to the SQLite database file + functions_data: List of dictionaries with function info extracted from DB + test_framework: 'pytest' or 'unittest' + max_run_count: Maximum number of runs to include in the test + + Returns: + A string containing the test code + + """ assert test_framework in ["pytest", "unittest"] imports = f"""import dill as pickle {"import unittest" if test_framework == "unittest" else ""} -from codeflash.tracing.replay_test import get_next_arg_and_return +from codeflash.benchmarking.replay_test import get_next_arg_and_return """ - # TODO: Module can have "-" character if the module-root is ".". Need to handle that case - function_properties: list[FunctionProperties] = [ - inspect_top_level_functions_or_methods( - file_name=function.file_name, - function_or_method_name=function.function_name, - class_name=function.class_name, - line_no=function.line_no, - ) - for function in functions - ] function_imports = [] - for function, function_property in zip(functions, function_properties): - if not function_property.is_top_level: - # can't be imported and run in the replay test - continue - if function_property.is_staticmethod: - function_imports.append( - f"from {function.module_name} import {function_property.staticmethod_class_name} as {get_function_alias(function.module_name, function_property.staticmethod_class_name)}" - ) - elif function.class_name: + for func in functions_data: + module_name = func.get("module_name") + function_name = func.get("function_name") + class_name = func.get("class_name", "") + + if class_name: function_imports.append( - f"from {function.module_name} import {function.class_name} as {get_function_alias(function.module_name, function.class_name)}" + f"from {module_name} import {class_name} as {get_function_alias(module_name, class_name)}" ) else: function_imports.append( - f"from {function.module_name} import {function.function_name} as {get_function_alias(function.module_name, function.function_name)}" + f"from {module_name} import {function_name} as {get_function_alias(module_name, function_name)}" ) imports += "\n".join(function_imports) - functions_to_optimize = [function.function_name for function in functions if function.function_name != "__init__"] + + functions_to_optimize = [func.get("function_name") for func in functions_data + if func.get("function_name") != "__init__"] metadata = f"""functions = {functions_to_optimize} trace_file_path = r"{trace_file}" -""" # trace_file_path path is parsed with regex later, format is important +""" + + # Templates for different types of tests test_function_body = textwrap.dedent( """\ - for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): - args = pickle.loads(arg_val_pkl) - ret = {function_name}({args}) + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + ret = {function_name}(*args, **kwargs) """ ) + test_class_method_body = textwrap.dedent( """\ - for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): - args = pickle.loads(arg_val_pkl){filter_variables} - ret = {class_name_alias}{method_name}(**args) - """ - ) - test_class_staticmethod_body = textwrap.dedent( - """\ - for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): - args = pickle.loads(arg_val_pkl){filter_variables} - ret = {class_name_alias}{method_name}(**args) + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl){filter_variables} + ret = {class_name_alias}{method_name}(**args, **kwargs) """ ) + if test_framework == "unittest": self = "self" test_template = "\nclass TestTracedFunctions(unittest.TestCase):\n" else: test_template = "" self = "" - for func, func_property in zip(functions, function_properties): - if not func_property.is_top_level: - # can't be imported and run in the replay test - continue - if func.class_name is None and not func_property.is_staticmethod: - alias = get_function_alias(func.module_name, func.function_name) + + for func in functions_data: + module_name = func.get("module_name") + function_name = func.get("function_name") + class_name = func.get("class_name") + file_name = func.get("file_name") + + if not class_name: + alias = get_function_alias(module_name, function_name) test_body = test_function_body.format( function_name=alias, - file_name=func.file_name, - orig_function_name=func.function_name, + file_name=file_name, + orig_function_name=function_name, max_run_count=max_run_count, - args="**args" if func_property.has_args else "", - ) - elif func_property.is_staticmethod: - class_name_alias = get_function_alias(func.module_name, func_property.staticmethod_class_name) - alias = get_function_alias( - func.module_name, func_property.staticmethod_class_name + "_" + func.function_name - ) - method_name = "." + func.function_name if func.function_name != "__init__" else "" - test_body = test_class_staticmethod_body.format( - orig_function_name=func.function_name, - file_name=func.file_name, - class_name_alias=class_name_alias, - method_name=method_name, - max_run_count=max_run_count, - filter_variables="", ) else: - class_name_alias = get_function_alias(func.module_name, func.class_name) - alias = get_function_alias(func.module_name, func.class_name + "_" + func.function_name) - - if func_property.is_classmethod: - filter_variables = '\n args.pop("cls", None)' - elif func.function_name == "__init__": - filter_variables = '\n args.pop("__class__", None)' - else: - filter_variables = "" - method_name = "." + func.function_name if func.function_name != "__init__" else "" + class_name_alias = get_function_alias(module_name, class_name) + alias = get_function_alias(module_name, class_name + "_" + function_name) + + filter_variables = "" + method_name = "." + function_name if function_name != "__init__" else "" test_body = test_class_method_body.format( - orig_function_name=func.function_name, - file_name=func.file_name, + orig_function_name=function_name, + file_name=file_name, class_name_alias=class_name_alias, - class_name=func.class_name, + class_name=class_name, method_name=method_name, max_run_count=max_run_count, filter_variables=filter_variables, ) + formatted_test_body = textwrap.indent(test_body, " " if test_framework == "unittest" else " ") test_template += " " if test_framework == "unittest" else "" diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 54e0b5118..5c0a077dc 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -3,12 +3,13 @@ from pathlib import Path import subprocess -def trace_benchmarks_pytest(benchmarks_root: Path, project_root: Path, output_file: Path) -> None: +def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, output_file: Path) -> None: result = subprocess.run( [ SAFE_SYS_EXECUTABLE, Path(__file__).parent / "pytest_new_process_trace_benchmarks.py", benchmarks_root, + tests_root, output_file, ], cwd=project_root, diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index df0c80fe2..fd03ab853 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -137,8 +137,10 @@ def optimize_function(self) -> Result[BestOptimization, str]: original_helper_code[helper_function_path] = helper_code if has_any_async_functions(code_context.read_writable_code): return Failure("Codeflash does not support async functions in the code to optimize.") - code_print(code_context.read_writable_code) + code_print(code_context.read_writable_code) + logger.info("Read only code") + code_print(code_context.read_only_context_code) generated_test_paths = [ get_test_file_path( self.test_cfg.tests_root, self.function_to_optimize.function_name, test_index, test_type="unit" diff --git a/codeflash/tracer.py b/codeflash/tracer.py index b63108677..eb4df84d4 100644 --- a/codeflash/tracer.py +++ b/codeflash/tracer.py @@ -18,19 +18,21 @@ import os import pathlib import pickle -import re import sqlite3 import sys +import threading import time +from argparse import ArgumentParser from collections import defaultdict -from copy import copy -from io import StringIO from pathlib import Path -from types import FrameType -from typing import Any, ClassVar, List +from typing import TYPE_CHECKING, Any, Callable, ClassVar import dill import isort +from rich.align import Align +from rich.panel import Panel +from rich.table import Table +from rich.text import Text from codeflash.cli_cmds.cli import project_root_from_module_root from codeflash.cli_cmds.console import console @@ -40,14 +42,34 @@ from codeflash.tracing.replay_test import create_trace_replay_test from codeflash.tracing.tracing_utils import FunctionModules from codeflash.verification.verification_utils import get_test_file_path -# import warnings -# warnings.filterwarnings("ignore", category=dill.PickleWarning) -# warnings.filterwarnings("ignore", category=DeprecationWarning) + +if TYPE_CHECKING: + from types import FrameType, TracebackType + + +class FakeCode: + def __init__(self, filename: str, line: int, name: str) -> None: + self.co_filename = filename + self.co_line = line + self.co_name = name + self.co_firstlineno = 0 + + def __repr__(self) -> str: + return repr((self.co_filename, self.co_line, self.co_name, None)) + + +class FakeFrame: + def __init__(self, code: FakeCode, prior: FakeFrame | None) -> None: + self.f_code = code + self.f_back = prior + self.f_locals: dict = {} + # Debug this file by simply adding print statements. This file is not meant to be debugged by the debugger. class Tracer: - """Use this class as a 'with' context manager to trace a function call, - input arguments, and profiling info. + """Use this class as a 'with' context manager to trace a function call. + + Traces function calls, input arguments, and profiling info. """ def __init__( @@ -58,9 +80,10 @@ def __init__( config_file_path: Path | None = None, max_function_count: int = 256, timeout: int | None = None, # seconds - benchmark: bool = False, ) -> None: - """:param output: The path to the output trace file + """Use this class to trace function calls. + + :param output: The path to the output trace file :param functions: List of functions to trace. If None, trace all functions :param disable: Disable the tracer if True :param config_file_path: Path to the pyproject.toml file, if None then it will be auto-discovered @@ -71,7 +94,9 @@ def __init__( if functions is None: functions = [] if os.environ.get("CODEFLASH_TRACER_DISABLE", "0") == "1": - console.print("Codeflash: Tracer disabled by environment variable CODEFLASH_TRACER_DISABLE") + console.rule( + "Codeflash: Tracer disabled by environment variable CODEFLASH_TRACER_DISABLE", style="bold red" + ) disable = True self.disable = disable if self.disable: @@ -96,17 +121,16 @@ def __init__( self.max_function_count = max_function_count self.config, found_config_path = parse_config_file(config_file_path) self.project_root = project_root_from_module_root(Path(self.config["module_root"]), found_config_path) - print("project_root", self.project_root) + console.rule(f"Project Root: {self.project_root}", style="bold blue") self.ignored_functions = {"", "", "", "", "", ""} - self.file_being_called_from: str = str(Path(sys._getframe().f_back.f_code.co_filename).name).replace(".", "_") + self.file_being_called_from: str = str(Path(sys._getframe().f_back.f_code.co_filename).name).replace(".", "_") # noqa: SLF001 assert timeout is None or timeout > 0, "Timeout should be greater than 0" self.timeout = timeout self.next_insert = 1000 self.trace_count = 0 - self.benchmark = benchmark # Profiler variables self.bias = 0 # calibration constant self.timings = {} @@ -121,48 +145,44 @@ def __init__( def __enter__(self) -> None: if self.disable: return - - # if getattr(Tracer, "used_once", False): - # console.print( - # "Codeflash: Tracer can only be used once per program run. " - # "Please only enable the Tracer once. Skipping tracing this section." - # ) - # self.disable = True - # return - # Tracer.used_once = True + if getattr(Tracer, "used_once", False): + console.print( + "Codeflash: Tracer can only be used once per program run. " + "Please only enable the Tracer once. Skipping tracing this section." + ) + self.disable = True + return + Tracer.used_once = True if pathlib.Path(self.output_file).exists(): - console.print("Codeflash: Removing existing trace file") + console.rule("Removing existing trace file", style="bold red") + console.rule() pathlib.Path(self.output_file).unlink(missing_ok=True) - self.con = sqlite3.connect(self.output_file) + self.con = sqlite3.connect(self.output_file, check_same_thread=False) cur = self.con.cursor() cur.execute("""PRAGMA synchronous = OFF""") + cur.execute("""PRAGMA journal_mode = WAL""") # TODO: Check out if we need to export the function test name as well cur.execute( "CREATE TABLE function_calls(type TEXT, function TEXT, classname TEXT, filename TEXT, " "line_number INTEGER, last_frame_address INTEGER, time_ns INTEGER, args BLOB)" ) - console.print("Codeflash: Tracing started!") - frame = sys._getframe(0) # Get this frame and simulate a call to it + console.rule("Codeflash: Traced Program Output Begin", style="bold blue") + frame = sys._getframe(0) # Get this frame and simulate a call to it # noqa: SLF001 self.dispatch["call"](self, frame, 0) self.start_time = time.time() sys.setprofile(self.trace_callback) + threading.setprofile(self.trace_callback) - def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + def __exit__( + self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None + ) -> None: if self.disable: return sys.setprofile(None) self.con.commit() - # Check if any functions were actually traced - if self.trace_count == 0: - self.con.close() - # Delete the trace file if no functions were traced - if self.output_file.exists(): - self.output_file.unlink() - console.print("Codeflash: No functions were traced. Removing trace database.") - return - + console.rule("Codeflash: Traced Program Output End", style="bold blue") self.create_stats() cur = self.con.cursor() @@ -186,25 +206,18 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: cur.execute("INSERT INTO total_time VALUES (?)", (self.total_tt,)) self.con.commit() self.con.close() - function_string = [str(function.file_name) + ":" + (function.class_name + ":" if function.class_name else "") + function.function_name for function in self.function_modules] - # print(function_string) # filter any functions where we did not capture the return - # self.function_modules = [ - # function - # for function in self.function_modules - # if self.function_count[ - # str(function.file_name) - # + ":" - # + (function.class_name + ":" if function.class_name else "") - # + function.function_name - # ] - # > 0 - # ] self.function_modules = [ function for function in self.function_modules - if str(str(function.file_name) + ":" + (function.class_name + ":" if function.class_name else "") + function.function_name) in self.function_count + if self.function_count[ + str(function.file_name) + + ":" + + (function.class_name + ":" if function.class_name else "") + + function.function_name + ] + > 0 ] replay_test = create_trace_replay_test( @@ -213,24 +226,15 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: test_framework=self.config["test_framework"], max_run_count=self.max_function_count, ) - # Need a better way to store the replay test - # function_path = "_".join(self.functions) if self.functions else self.file_being_called_from - function_path = self.file_being_called_from - if self.benchmark and self.config["benchmarks_root"]: - # check if replay test dir exists, create - replay_test_dir = Path(self.config["benchmarks_root"]) / "codeflash_replay_tests" - if not replay_test_dir.exists(): - replay_test_dir.mkdir(parents=True) - test_file_path = get_test_file_path( - test_dir=replay_test_dir, function_name=function_path, test_type="replay" - ) - else: - test_file_path = get_test_file_path( - test_dir=Path(self.config["tests_root"]), function_name=function_path, test_type="replay" - ) + function_path = "_".join(self.functions) if self.functions else self.file_being_called_from + test_file_path = get_test_file_path( + test_dir=Path(self.config["tests_root"]), function_name=function_path, test_type="replay" + ) replay_test = isort.code(replay_test) - with open(test_file_path, "w", encoding="utf8") as file: + + with Path(test_file_path).open("w", encoding="utf8") as file: file.write(replay_test) + console.print( f"Codeflash: Traced {self.trace_count} function calls successfully and replay test created at - {test_file_path}", crop=False, @@ -253,9 +257,8 @@ def tracer_logic(self, frame: FrameType, event: str) -> None: return if not file_name.exists(): return - # if self.functions: - # if code.co_name not in self.functions: - # return + if self.functions and code.co_name not in self.functions: + return class_name = None arguments = frame.f_locals try: @@ -267,16 +270,12 @@ def tracer_logic(self, frame: FrameType, event: str) -> None: class_name = arguments["self"].__class__.__name__ elif "cls" in arguments and hasattr(arguments["cls"], "__name__"): class_name = arguments["cls"].__name__ - except: + except: # noqa: E722 # someone can override the getattr method and raise an exception. I'm looking at you wrapt return - function_qualified_name = f"{file_name}:{(class_name + ':' if class_name else '')}{code.co_name}" if function_qualified_name in self.ignored_qualified_functions: return - if self.functions and function_qualified_name not in self.functions: - return - if function_qualified_name not in self.function_count: # seeing this function for the first time self.function_count[function_qualified_name] = 0 @@ -351,7 +350,7 @@ def tracer_logic(self, frame: FrameType, event: str) -> None: self.next_insert = 1000 self.con.commit() - def trace_callback(self, frame: FrameType, event: str, arg: Any) -> None: + def trace_callback(self, frame: FrameType, event: str, arg: str | None) -> None: # profiler section timer = self.timer t = timer() - self.t - self.bias @@ -367,45 +366,60 @@ def trace_callback(self, frame: FrameType, event: str, arg: Any) -> None: else: self.t = timer() - t # put back unrecorded delta - def trace_dispatch_call(self, frame, t) -> int: - if self.cur and frame.f_back is not self.cur[-2]: - rpt, rit, ret, rfn, rframe, rcur = self.cur - if not isinstance(rframe, Tracer.fake_frame): - assert rframe.f_back is frame.f_back, ("Bad call", rfn, rframe, rframe.f_back, frame, frame.f_back) - self.trace_dispatch_return(rframe, 0) - assert self.cur is None or frame.f_back is self.cur[-2], ("Bad call", self.cur[-3]) - fcode = frame.f_code - arguments = frame.f_locals - class_name = None + def trace_dispatch_call(self, frame: FrameType, t: int) -> int: + """Handle call events in the profiler.""" try: - if ( - "self" in arguments - and hasattr(arguments["self"], "__class__") - and hasattr(arguments["self"].__class__, "__name__") - ): - class_name = arguments["self"].__class__.__name__ - elif "cls" in arguments and hasattr(arguments["cls"], "__name__"): - class_name = arguments["cls"].__name__ - except: - pass - fn = (fcode.co_filename, fcode.co_firstlineno, fcode.co_name, class_name) - self.cur = (t, 0, 0, fn, frame, self.cur) - timings = self.timings - if fn in timings: - cc, ns, tt, ct, callers = timings[fn] - timings[fn] = cc, ns + 1, tt, ct, callers - else: - timings[fn] = 0, 0, 0, 0, {} - return 1 + # In multi-threaded contexts, we need to be more careful about frame comparisons + if self.cur and frame.f_back is not self.cur[-2]: + # This happens when we're in a different thread + rpt, rit, ret, rfn, rframe, rcur = self.cur + + # Only attempt to handle the frame mismatch if we have a valid rframe + if ( + not isinstance(rframe, FakeFrame) + and hasattr(rframe, "f_back") + and hasattr(frame, "f_back") + and rframe.f_back is frame.f_back + ): + self.trace_dispatch_return(rframe, 0) + + # Get function information + fcode = frame.f_code + arguments = frame.f_locals + class_name = None + try: + if ( + "self" in arguments + and hasattr(arguments["self"], "__class__") + and hasattr(arguments["self"].__class__, "__name__") + ): + class_name = arguments["self"].__class__.__name__ + elif "cls" in arguments and hasattr(arguments["cls"], "__name__"): + class_name = arguments["cls"].__name__ + except Exception: # noqa: BLE001, S110 + pass + + fn = (fcode.co_filename, fcode.co_firstlineno, fcode.co_name, class_name) + self.cur = (t, 0, 0, fn, frame, self.cur) + timings = self.timings + if fn in timings: + cc, ns, tt, ct, callers = timings[fn] + timings[fn] = cc, ns + 1, tt, ct, callers + else: + timings[fn] = 0, 0, 0, 0, {} + return 1 # noqa: TRY300 + except Exception: # noqa: BLE001 + # Handle any errors gracefully + return 0 - def trace_dispatch_exception(self, frame, t): + def trace_dispatch_exception(self, frame: FrameType, t: int) -> int: rpt, rit, ret, rfn, rframe, rcur = self.cur if (rframe is not frame) and rcur: return self.trace_dispatch_return(rframe, t) self.cur = rpt, rit + t, ret, rfn, rframe, rcur return 1 - def trace_dispatch_c_call(self, frame, t) -> int: + def trace_dispatch_c_call(self, frame: FrameType, t: int) -> int: fn = ("", 0, self.c_func_name, None) self.cur = (t, 0, 0, fn, frame, self.cur) timings = self.timings @@ -416,15 +430,27 @@ def trace_dispatch_c_call(self, frame, t) -> int: timings[fn] = 0, 0, 0, 0, {} return 1 - def trace_dispatch_return(self, frame, t) -> int: - if frame is not self.cur[-2]: - assert frame is self.cur[-2].f_back, ("Bad return", self.cur[-3]) - self.trace_dispatch_return(self.cur[-2], 0) + def trace_dispatch_return(self, frame: FrameType, t: int) -> int: + if not self.cur or not self.cur[-2]: + return 0 + # In multi-threaded environments, frames can get mismatched + if frame is not self.cur[-2]: + # Don't assert in threaded environments - frames can legitimately differ + if hasattr(frame, "f_back") and hasattr(self.cur[-2], "f_back") and frame is self.cur[-2].f_back: + self.trace_dispatch_return(self.cur[-2], 0) + else: + # We're in a different thread or context, can't continue with this frame + return 0 # Prefix "r" means part of the Returning or exiting frame. # Prefix "p" means part of the Previous or Parent or older frame. rpt, rit, ret, rfn, frame, rcur = self.cur + + # Guard against invalid rcur (w threading) + if not rcur: + return 0 + rit = rit + t frame_total = rit + ret @@ -432,6 +458,9 @@ def trace_dispatch_return(self, frame, t) -> int: self.cur = ppt, pit + rpt, pet + frame_total, pfn, pframe, pcur timings = self.timings + if rfn not in timings: + # w threading, rfn can be missing + timings[rfn] = 0, 0, 0, 0, {} cc, ns, tt, ct, callers = timings[rfn] if not ns: # This is the only occurrence of the function on the stack. @@ -453,7 +482,7 @@ def trace_dispatch_return(self, frame, t) -> int: return 1 - dispatch: ClassVar[dict[str, callable]] = { + dispatch: ClassVar[dict[str, Callable[[Tracer, FrameType, int], int]]] = { "call": trace_dispatch_call, "exception": trace_dispatch_exception, "return": trace_dispatch_return, @@ -462,26 +491,10 @@ def trace_dispatch_return(self, frame, t) -> int: "c_return": trace_dispatch_return, } - class fake_code: - def __init__(self, filename, line, name) -> None: - self.co_filename = filename - self.co_line = line - self.co_name = name - self.co_firstlineno = 0 - - def __repr__(self) -> str: - return repr((self.co_filename, self.co_line, self.co_name, None)) - - class fake_frame: - def __init__(self, code, prior) -> None: - self.f_code = code - self.f_back = prior - self.f_locals = {} - - def simulate_call(self, name) -> None: - code = self.fake_code("profiler", 0, name) + def simulate_call(self, name: str) -> None: + code = FakeCode("profiler", 0, name) pframe = self.cur[-2] if self.cur else None - frame = self.fake_frame(code, pframe) + frame = FakeFrame(code, pframe) self.dispatch["call"](self, frame, 0) def simulate_cmd_complete(self) -> None: @@ -494,58 +507,172 @@ def simulate_cmd_complete(self) -> None: t = 0 self.t = get_time() - t - def print_stats(self, sort=-1) -> None: - import pstats + def print_stats(self, sort: str | int | tuple = -1) -> None: + if not self.stats: + console.print("Codeflash: No stats available to print") + self.total_tt = 0 + return if not isinstance(sort, tuple): sort = (sort,) - # The following code customizes the default printing behavior to - # print in milliseconds. - s = StringIO() - stats_obj = pstats.Stats(copy(self), stream=s) - stats_obj.strip_dirs().sort_stats(*sort).print_stats(100) - self.total_tt = stats_obj.total_tt - console.print("total_tt", self.total_tt) - raw_stats = s.getvalue() - m = re.search(r"function calls?.*in (\d+)\.\d+ (seconds?)", raw_stats) - total_time = None - if m: - total_time = int(m.group(1)) - if total_time is None: - console.print("Failed to get total time from stats") - total_time_ms = total_time / 1e6 - raw_stats = re.sub( - r"(function calls?.*)in (\d+)\.\d+ (seconds?)", rf"\1 in {total_time_ms:.3f} milliseconds", raw_stats - ) - match_pattern = r"^ *[\d\/]+ +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +" - m = re.findall(match_pattern, raw_stats, re.MULTILINE) - ms_times = [] - for tottime, percall, cumtime, percall_cum in m: - tottime_ms = int(tottime) / 1e6 - percall_ms = int(percall) / 1e6 - cumtime_ms = int(cumtime) / 1e6 - percall_cum_ms = int(percall_cum) / 1e6 - ms_times.append([tottime_ms, percall_ms, cumtime_ms, percall_cum_ms]) - split_stats = raw_stats.split("\n") - new_stats = [] - - replace_pattern = r"^( *[\d\/]+) +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +(.*)" - times_index = 0 - for line in split_stats: - if times_index >= len(ms_times): - replaced = line - else: - replaced, n = re.subn( - replace_pattern, - rf"\g<1>{ms_times[times_index][0]:8.3f} {ms_times[times_index][1]:8.3f} {ms_times[times_index][2]:8.3f} {ms_times[times_index][3]:8.3f} \g<6>", - line, - count=1, + + # First, convert stats to make them pstats-compatible + try: + # Initialize empty collections for pstats + self.files = [] + self.top_level = [] + + # Create entirely new dictionaries instead of modifying existing ones + new_stats = {} + new_timings = {} + + # Convert stats dictionary + stats_items = list(self.stats.items()) + for func, stats_data in stats_items: + try: + # Make sure we have 5 elements in stats_data + if len(stats_data) != 5: + console.print(f"Skipping malformed stats data for {func}: {stats_data}") + continue + + cc, nc, tt, ct, callers = stats_data + + if len(func) == 4: + file_name, line_num, func_name, class_name = func + new_func_name = f"{class_name}.{func_name}" if class_name else func_name + new_func = (file_name, line_num, new_func_name) + else: + new_func = func # Keep as is if already in correct format + + new_callers = {} + callers_items = list(callers.items()) + for caller_func, count in callers_items: + if isinstance(caller_func, tuple): + if len(caller_func) == 4: + caller_file, caller_line, caller_name, caller_class = caller_func + caller_new_name = f"{caller_class}.{caller_name}" if caller_class else caller_name + new_caller_func = (caller_file, caller_line, caller_new_name) + else: + new_caller_func = caller_func + else: + console.print(f"Unexpected caller format: {caller_func}") + new_caller_func = str(caller_func) + + new_callers[new_caller_func] = count + + # Store with new format + new_stats[new_func] = (cc, nc, tt, ct, new_callers) + except Exception as e: # noqa: BLE001 + console.print(f"Error converting stats for {func}: {e}") + continue + + timings_items = list(self.timings.items()) + for func, timing_data in timings_items: + try: + if len(timing_data) != 5: + console.print(f"Skipping malformed timing data for {func}: {timing_data}") + continue + + cc, ns, tt, ct, callers = timing_data + + if len(func) == 4: + file_name, line_num, func_name, class_name = func + new_func_name = f"{class_name}.{func_name}" if class_name else func_name + new_func = (file_name, line_num, new_func_name) + else: + new_func = func + + new_callers = {} + callers_items = list(callers.items()) + for caller_func, count in callers_items: + if isinstance(caller_func, tuple): + if len(caller_func) == 4: + caller_file, caller_line, caller_name, caller_class = caller_func + caller_new_name = f"{caller_class}.{caller_name}" if caller_class else caller_name + new_caller_func = (caller_file, caller_line, caller_new_name) + else: + new_caller_func = caller_func + else: + console.print(f"Unexpected caller format: {caller_func}") + new_caller_func = str(caller_func) + + new_callers[new_caller_func] = count + + new_timings[new_func] = (cc, ns, tt, ct, new_callers) + except Exception as e: # noqa: BLE001 + console.print(f"Error converting timings for {func}: {e}") + continue + + self.stats = new_stats + self.timings = new_timings + + self.total_tt = sum(tt for _, _, tt, _, _ in self.stats.values()) + + total_calls = sum(cc for cc, _, _, _, _ in self.stats.values()) + total_primitive = sum(nc for _, nc, _, _, _ in self.stats.values()) + + summary = Text.assemble( + f"{total_calls:,} function calls ", + ("(" + f"{total_primitive:,} primitive calls" + ")", "dim"), + f" in {self.total_tt / 1e6:.3f}milliseconds", + ) + + console.print(Align.center(Panel(summary, border_style="blue", width=80, padding=(0, 2), expand=False))) + + table = Table( + show_header=True, + header_style="bold magenta", + border_style="blue", + title="[bold]Function Profile[/bold] (ordered by internal time)", + title_style="cyan", + caption=f"Showing top 25 of {len(self.stats)} functions", + ) + + table.add_column("Calls", justify="right", style="green", width=10) + table.add_column("Time (ms)", justify="right", style="cyan", width=10) + table.add_column("Per Call", justify="right", style="cyan", width=10) + table.add_column("Cum (ms)", justify="right", style="yellow", width=10) + table.add_column("Cum/Call", justify="right", style="yellow", width=10) + table.add_column("Function", style="blue") + + sorted_stats = sorted( + ((func, stats) for func, stats in self.stats.items() if isinstance(func, tuple) and len(func) == 3), + key=lambda x: x[1][2], # Sort by tt (internal time) + reverse=True, + )[:25] # Limit to top 25 + + # Format and add each row to the table + for func, (cc, nc, tt, ct, _) in sorted_stats: + filename, lineno, funcname = func + + # Format calls - show recursive format if different + calls_str = f"{cc}/{nc}" if cc != nc else f"{cc:,}" + + # Convert to milliseconds + tt_ms = tt / 1e6 + ct_ms = ct / 1e6 + + # Calculate per-call times + per_call = tt_ms / cc if cc > 0 else 0 + cum_per_call = ct_ms / nc if nc > 0 else 0 + base_filename = Path(filename).name + file_link = f"[link=file://{filename}]{base_filename}[/link]" + + table.add_row( + calls_str, + f"{tt_ms:.3f}", + f"{per_call:.3f}", + f"{ct_ms:.3f}", + f"{cum_per_call:.3f}", + f"{funcname} [dim]({file_link}:{lineno})[/dim]", ) - if n > 0: - times_index += 1 - new_stats.append(replaced) - console.print("\n".join(new_stats)) + console.print(Align.center(table)) + + except Exception as e: # noqa: BLE001 + console.print(f"[bold red]Error in stats processing:[/bold red] {e}") + console.print(f"Traced {self.trace_count:,} function calls") + self.total_tt = 0 def make_pstats_compatible(self) -> None: # delete the extra class_name item from the function tuple @@ -562,9 +689,8 @@ def make_pstats_compatible(self) -> None: self.stats = new_stats self.timings = new_timings - def dump_stats(self, file) -> None: - with open(file, "wb") as f: - self.create_stats() + def dump_stats(self, file: str) -> None: + with Path(file).open("wb") as f: marshal.dump(self.stats, f) def create_stats(self) -> None: @@ -573,25 +699,23 @@ def create_stats(self) -> None: def snapshot_stats(self) -> None: self.stats = {} - for func, (cc, _ns, tt, ct, callers) in self.timings.items(): - callers = callers.copy() + for func, (cc, _ns, tt, ct, caller_dict) in self.timings.items(): + callers = caller_dict.copy() nc = 0 for callcnt in callers.values(): nc += callcnt self.stats[func] = cc, nc, tt, ct, callers - def runctx(self, cmd, globals, locals): + def runctx(self, cmd: str, global_vars: dict[str, Any], local_vars: dict[str, Any]) -> Tracer | None: self.__enter__() try: - exec(cmd, globals, locals) + exec(cmd, global_vars, local_vars) # noqa: S102 finally: self.__exit__(None, None, None) return self -def main(): - from argparse import ArgumentParser - +def main() -> ArgumentParser: parser = ArgumentParser(allow_abbrev=False) parser.add_argument("-o", "--outfile", dest="outfile", help="Save trace to ", required=True) parser.add_argument("--only-functions", help="Trace only these functions", nargs="+", default=None) @@ -648,16 +772,13 @@ def main(): "__cached__": None, } try: - tracer = Tracer( + Tracer( output=args.outfile, functions=args.only_functions, max_function_count=args.max_function_count, timeout=args.tracer_timeout, config_file_path=args.codeflash_config, - ) - - tracer.runctx(code, globs, None) - print(tracer.functions) + ).runctx(code, globs, None) except BrokenPipeError as exc: # Prevent "Exception ignored" during interpreter shutdown. diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 071535c6a..570888fcc 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,12 +1,27 @@ +from codeflash.benchmarking.codeflash_trace import codeflash_trace from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from pathlib import Path from codeflash.code_utils.code_utils import get_run_tmp_file +import shutil 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" - output_file = Path("test_trace_benchmarks.trace").resolve() - trace_benchmarks_pytest(benchmarks_root, project_root, output_file) + # make directory in project_root / "tests" + + + tests_root = project_root / "tests" / "test_trace_benchmarks" + tests_root.mkdir(parents=False, exist_ok=False) + output_file = (tests_root / Path("test_trace_benchmarks.trace")).resolve() + trace_benchmarks_pytest(benchmarks_root, tests_root, project_root, output_file) assert output_file.exists() - output_file.unlink() \ No newline at end of file + + test1_path = tests_root / Path("test_benchmark_bubble_sort_py_test_sort__replay_test_0.py") + assert test1_path.exists() + + # test1_code = """""" + # assert test1_path.read_text("utf-8").strip()==test1_code.strip() + # cleanup + # shutil.rmtree(tests_root) + # output_file.unlink() \ No newline at end of file From f7466a577b857ebdeef20ac43f58f2dd857b556e Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 14 Mar 2025 15:03:36 -0700 Subject: [PATCH 009/122] replay test functionality working for functions, methods, static methods, class methods, init. basic instrumentation logic for codeflash_trace done. --- .../bubble_sort_codeflash_trace.py | 46 ++ .../benchmarks/test_benchmark_bubble_sort.py | 17 +- .../benchmarks/test_process_and_sort.py | 6 +- codeflash/benchmarking/codeflash_trace.py | 106 +---- .../instrument_codeflash_trace.py | 79 ++++ codeflash/benchmarking/plugin/plugin.py | 6 + .../pytest_new_process_trace_benchmarks.py | 2 +- codeflash/benchmarking/replay_test.py | 170 ++++++- tests/test_instrument_codeflash_capture.py | 441 ++++++------------ tests/test_instrument_codeflash_trace.py | 239 ++++++++++ tests/test_trace_benchmarks.py | 147 +++++- 11 files changed, 829 insertions(+), 430 deletions(-) create mode 100644 code_to_optimize/bubble_sort_codeflash_trace.py create mode 100644 tests/test_instrument_codeflash_trace.py diff --git a/code_to_optimize/bubble_sort_codeflash_trace.py b/code_to_optimize/bubble_sort_codeflash_trace.py new file mode 100644 index 000000000..ee4dbd999 --- /dev/null +++ b/code_to_optimize/bubble_sort_codeflash_trace.py @@ -0,0 +1,46 @@ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +@codeflash_trace +def sorter(arr): + for i in range(len(arr)): + for j in range(len(arr) - 1): + if arr[j] > arr[j + 1]: + temp = arr[j] + arr[j] = arr[j + 1] + arr[j + 1] = temp + return arr + +class Sorter: + @codeflash_trace + def __init__(self, arr): + self.arr = arr + @codeflash_trace + def sorter(self, multiplier): + for i in range(len(self.arr)): + for j in range(len(self.arr) - 1): + if self.arr[j] > self.arr[j + 1]: + temp = self.arr[j] + self.arr[j] = self.arr[j + 1] + self.arr[j + 1] = temp + return self.arr * multiplier + + @staticmethod + @codeflash_trace + def sort_static(arr): + for i in range(len(arr)): + for j in range(len(arr) - 1): + if arr[j] > arr[j + 1]: + temp = arr[j] + arr[j] = arr[j + 1] + arr[j + 1] = temp + return arr + + @classmethod + @codeflash_trace + def sort_class(cls, arr): + for i in range(len(arr)): + for j in range(len(arr) - 1): + if arr[j] > arr[j + 1]: + temp = arr[j] + arr[j] = arr[j + 1] + arr[j + 1] = temp + return arr diff --git a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py index 21c2bbb29..03b9d38d1 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py @@ -1,13 +1,20 @@ import pytest -from code_to_optimize.bubble_sort_codeflash_trace import sorter +from code_to_optimize.bubble_sort_codeflash_trace import sorter, Sorter def test_sort(benchmark): - result = benchmark(sorter, list(reversed(range(5000)))) - assert result == list(range(5000)) + result = benchmark(sorter, list(reversed(range(500)))) + assert result == list(range(500)) # This should not be picked up as a benchmark test def test_sort2(): - result = sorter(list(reversed(range(5000)))) - assert result == list(range(5000)) \ No newline at end of file + result = sorter(list(reversed(range(500)))) + assert result == list(range(500)) + +def test_class_sort(benchmark): + obj = Sorter(list(reversed(range(100)))) + result1 = benchmark(obj.sorter, 2) + result2 = benchmark(Sorter.sort_class, list(reversed(range(100)))) + result3 = benchmark(Sorter.sort_static, list(reversed(range(100)))) + result4 = benchmark(Sorter, [1,2,3]) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py index 2713721e4..bcd42eab9 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py @@ -1,8 +1,8 @@ from code_to_optimize.process_and_bubble_sort_codeflash_trace import compute_and_sort from code_to_optimize.bubble_sort_codeflash_trace import sorter def test_compute_and_sort(benchmark): - result = benchmark(compute_and_sort, list(reversed(range(5000)))) - assert result == 6247083.5 + result = benchmark(compute_and_sort, list(reversed(range(500)))) + assert result == 62208.5 def test_no_func(benchmark): - benchmark(sorter, list(reversed(range(5000)))) \ No newline at end of file + benchmark(sorter, list(reversed(range(500)))) \ No newline at end of file diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index c678b7643..65ba98783 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -51,6 +51,10 @@ def wrapper(*args, **kwargs): overhead_time = 0 try: + # Check if currently in pytest benchmark fixture + if os.environ.get("CODEFLASH_BENCHMARKING", "False") == "False": + return result + # Pickle the arguments pickled_args = pickle.dumps(args, protocol=pickle.HIGHEST_PROTOCOL) pickled_kwargs = pickle.dumps(kwargs, protocol=pickle.HIGHEST_PROTOCOL) @@ -58,6 +62,7 @@ def wrapper(*args, **kwargs): # 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", "") + benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "") # Calculate overhead time overhead_end_time = time.time() @@ -69,7 +74,7 @@ def wrapper(*args, **kwargs): class_name = qualname.split(".")[0] self.function_calls_data.append( (func.__name__, class_name, func.__module__, func.__code__.co_filename, - benchmark_function_name, benchmark_file_name, execution_time, + benchmark_function_name, benchmark_file_name, benchmark_line_number, execution_time, overhead_time, pickled_args, pickled_kwargs) ) @@ -100,7 +105,7 @@ def write_to_db(self, output_file: str) -> None: cur.execute( "CREATE TABLE IF NOT EXISTS function_calls(" "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," - "benchmark_function_name TEXT, benchmark_file_name TEXT, " + "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," "time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" ) @@ -108,8 +113,8 @@ def write_to_db(self, output_file: str) -> None: cur.executemany( "INSERT INTO function_calls " "(function_name, class_name, module_name, file_name, benchmark_function_name, " - "benchmark_file_name, time_ns, overhead_time_ns, args, kwargs) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + "benchmark_file_name, benchmark_line_number, time_ns, overhead_time_ns, args, kwargs) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", self.function_calls_data ) @@ -165,13 +170,14 @@ def print_codeflash_db(self, limit: int = None) -> None: print(f" File: {row[3]}") print(f" Benchmark Function: {row[4] or 'N/A'}") print(f" Benchmark File: {row[5] or 'N/A'}") - print(f" Execution Time: {row[6]:.6f} seconds") - print(f" Overhead Time: {row[7]:.6f} seconds") + print(f" Benchmark Line: {row[6] or 'N/A'}") + print(f" Execution Time: {row[7]:.6f} seconds") + print(f" Overhead Time: {row[8]:.6f} seconds") # Unpickle and print args and kwargs try: - args = pickle.loads(row[8]) - kwargs = pickle.loads(row[9]) + args = pickle.loads(row[9]) + kwargs = pickle.loads(row[10]) print(f" Args: {args}") print(f" Kwargs: {kwargs}") @@ -187,90 +193,6 @@ def print_codeflash_db(self, limit: int = None) -> None: except Exception as e: print(f"Error reading database: {e}") - def generate_replay_test(self, output_dir: str = None, project_root: str = "", test_framework: str = "pytest", - max_run_count: int = 100) -> None: - """ - Generate multiple replay tests from the traced function calls, grouping by benchmark name. - - Args: - output_dir: Directory to write the generated tests (if None, only returns the code) - project_root: Root directory of the project for module imports - test_framework: 'pytest' or 'unittest' - max_run_count: Maximum number of runs to include per function - - Returns: - Dictionary mapping benchmark names to generated test code - """ - import isort - from codeflash.verification.verification_utils import get_test_file_path - - if not self.db_path: - print("No database path set. Call write_to_db first or set db_path manually.") - return {} - - try: - # Import the function here to avoid circular imports - from codeflash.benchmarking.replay_test import create_trace_replay_test - - print("connecting to: ", self.db_path) - # Connect to the database - conn = sqlite3.connect(self.db_path) - cursor = conn.cursor() - - # Get distinct benchmark names - cursor.execute( - "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM function_calls" - ) - benchmarks = cursor.fetchall() - - # Generate a test for each benchmark - for benchmark in benchmarks: - benchmark_function_name, benchmark_file_name = benchmark - # Get functions associated with this benchmark - cursor.execute( - "SELECT DISTINCT function_name, class_name, module_name, file_name FROM function_calls " - "WHERE benchmark_function_name = ? AND benchmark_file_name = ?", - (benchmark_function_name, benchmark_file_name) - ) - - functions_data = [] - for func_row in cursor.fetchall(): - function_name, class_name, module_name, file_name = func_row - - # Add this function to our list - functions_data.append({ - "function_name": function_name, - "class_name": class_name, - "file_name": file_name, - "module_name": module_name - }) - - if not functions_data: - print(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") - continue - - # Generate the test code for this benchmark - test_code = create_trace_replay_test( - trace_file=self.db_path, - functions_data=functions_data, - test_framework=test_framework, - max_run_count=max_run_count, - ) - test_code = isort.code(test_code) - - # Write to file if requested - if output_dir: - output_file = get_test_file_path( - test_dir=Path(output_dir), function_name=f"{benchmark_file_name[5:]}_{benchmark_function_name}", test_type="replay" - ) - with open(output_file, 'w') as f: - f.write(test_code) - print(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") - - conn.close() - - except Exception as e: - print(f"Error generating replay tests: {e}") # Create a singleton instance codeflash_trace = CodeflashTrace() diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py index 8b1378917..99b2dad20 100644 --- a/codeflash/benchmarking/instrument_codeflash_trace.py +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -1 +1,80 @@ +import libcst as cst +from codeflash.discovery.functions_to_optimize import FunctionToOptimize + + +class AddDecoratorTransformer(cst.CSTTransformer): + def __init__(self, function_name, class_name=None): + super().__init__() + self.function_name = function_name + self.class_name = class_name + self.in_target_class = (class_name is None) # If no class name, always "in target class" + + def leave_ClassDef(self, original_node, updated_node): + if self.class_name and original_node.name.value == self.class_name: + self.in_target_class = False + return updated_node + + def visit_ClassDef(self, node): + if self.class_name and node.name.value == self.class_name: + self.in_target_class = True + return True + + def leave_FunctionDef(self, original_node, updated_node): + if not self.in_target_class or original_node.name.value != self.function_name: + return updated_node + + # Create the codeflash_trace decorator + decorator = cst.Decorator( + decorator=cst.Name(value="codeflash_trace") + ) + + # Add the new decorator after any existing decorators + updated_decorators = list(updated_node.decorators) + [decorator] + + # Return the updated node with the new decorator + return updated_node.with_changes( + decorators=updated_decorators + ) + + +def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToOptimize) -> str: + """Add codeflash_trace to a function. + + Args: + code: The source code as a string + function_to_optimize: The FunctionToOptimize instance containing function details + + Returns: + The modified source code as a string + """ + # Extract class name if present + class_name = None + if len(function_to_optimize.parents) == 1 and function_to_optimize.parents[0].type == "ClassDef": + class_name = function_to_optimize.parents[0].name + + transformer = AddDecoratorTransformer( + function_name=function_to_optimize.function_name, + class_name=class_name + ) + + module = cst.parse_module(code) + modified_module = module.visit(transformer) + return modified_module.code + + +def instrument_codeflash_trace( + function_to_optimize: FunctionToOptimize +) -> None: + """Instrument __init__ function with codeflash_trace decorator if it's in a class.""" + # Instrument fto class + original_code = function_to_optimize.file_path.read_text(encoding="utf-8") + + # Modify the code + modified_code = add_codeflash_decorator_to_code( + original_code, + function_to_optimize + ) + + # Write the modified code back to the file + function_to_optimize.file_path.write_text(modified_code, encoding="utf-8") diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index 6d8db9bf9..caf175a4e 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -1,3 +1,5 @@ +import sys + import pytest import time import os @@ -34,12 +36,16 @@ def benchmark(request): return None class Benchmark: + def __call__(self, func, *args, **kwargs): os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = request.node.name os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = request.node.fspath.basename + os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack + os.environ["CODEFLASH_BENCHMARKING"] = "True" start = time.process_time_ns() result = func(*args, **kwargs) end = time.process_time_ns() + os.environ["CODEFLASH_BENCHMARKING"] = "False" print(f"Benchmark: {func.__name__} took {end - start} ns") return result diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 8e1958fec..04c5e67ea 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -20,7 +20,7 @@ ) codeflash_trace.write_to_db(output_file) codeflash_trace.print_codeflash_db() - codeflash_trace.generate_replay_test(tests_root, project_root, test_framework="pytest") + except Exception as e: print(f"Failed to collect tests: {e!s}") exitcode = -1 \ No newline at end of file diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 9bc2c79f3..58ce456c2 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -5,6 +5,12 @@ from collections.abc import Generator from typing import Any, Dict +import isort + +from codeflash.cli_cmds.console import logger +from codeflash.discovery.functions_to_optimize import inspect_top_level_functions_or_methods +from codeflash.verification.verification_utils import get_test_file_path +from pathlib import Path def get_next_arg_and_return( trace_file: str, function_name: str, file_name: str, class_name: str | None = None, num_to_get: int = 25 @@ -20,21 +26,21 @@ def get_next_arg_and_return( ) else: cursor = cur.execute( - "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? ORDER BY time_ns ASC LIMIT ?", + "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? AND class_name = '' ORDER BY time_ns ASC LIMIT ?", (function_name, file_name, limit), ) while (val := cursor.fetchone()) is not None: - yield val[8], val[9] # args and kwargs are at indices 7 and 8 + yield val[9], val[10] # args and kwargs are at indices 7 and 8 def get_function_alias(module: str, function_name: str) -> str: return "_".join(module.split(".")) + "_" + function_name -def create_trace_replay_test( +def create_trace_replay_test_code( trace_file: str, - functions_data: list[Dict[str, Any]], + functions_data: list[dict[str, Any]], test_framework: str = "pytest", max_run_count=100 ) -> str: @@ -52,7 +58,7 @@ def create_trace_replay_test( """ assert test_framework in ["pytest", "unittest"] - imports = f"""import dill as pickle + imports = f"""import dill as pickle {"import unittest" if test_framework == "unittest" else ""} from codeflash.benchmarking.replay_test import get_next_arg_and_return """ @@ -62,7 +68,6 @@ def create_trace_replay_test( module_name = func.get("module_name") function_name = func.get("function_name") class_name = func.get("class_name", "") - if class_name: function_imports.append( f"from {module_name} import {class_name} as {get_function_alias(module_name, class_name)}" @@ -90,12 +95,37 @@ def create_trace_replay_test( """ ) + test_method_body = textwrap.dedent( + """\ + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl){filter_variables} + function_name = "{orig_function_name}" + if not args: + raise ValueError("No arguments provided for the method.") + if function_name == "__init__": + ret = {class_name_alias}(*args[1:], **kwargs) + else: + instance = args[0] # self + ret = instance{method_name}(*args[1:], **kwargs) + """) + test_class_method_body = textwrap.dedent( """\ for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} - ret = {class_name_alias}{method_name}(**args, **kwargs) + if not args: + raise ValueError("No arguments provided for the method.") + ret = {class_name_alias}{method_name}(*args[1:], **kwargs) + """ + ) + test_static_method_body = textwrap.dedent( + """\ + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl){filter_variables} + ret = {class_name_alias}{method_name}(*args, **kwargs) """ ) @@ -111,7 +141,9 @@ def create_trace_replay_test( function_name = func.get("function_name") class_name = func.get("class_name") file_name = func.get("file_name") - + function_properties = func.get("function_properties") + print(f"Class: {class_name}, Function: {function_name}") + print(function_properties) if not class_name: alias = get_function_alias(module_name, function_name) test_body = test_function_body.format( @@ -125,16 +157,38 @@ def create_trace_replay_test( alias = get_function_alias(module_name, class_name + "_" + function_name) filter_variables = "" + # filter_variables = '\n args.pop("cls", None)' method_name = "." + function_name if function_name != "__init__" else "" - test_body = test_class_method_body.format( - orig_function_name=function_name, - file_name=file_name, - class_name_alias=class_name_alias, - class_name=class_name, - method_name=method_name, - max_run_count=max_run_count, - filter_variables=filter_variables, - ) + if function_properties.is_classmethod: + test_body = test_class_method_body.format( + orig_function_name=function_name, + file_name=file_name, + class_name_alias=class_name_alias, + class_name=class_name, + method_name=method_name, + max_run_count=max_run_count, + filter_variables=filter_variables, + ) + elif function_properties.is_staticmethod: + test_body = test_static_method_body.format( + orig_function_name=function_name, + file_name=file_name, + class_name_alias=class_name_alias, + class_name=class_name, + method_name=method_name, + max_run_count=max_run_count, + filter_variables=filter_variables, + ) + else: + test_body = test_method_body.format( + orig_function_name=function_name, + file_name=file_name, + class_name_alias=class_name_alias, + class_name=class_name, + method_name=method_name, + max_run_count=max_run_count, + filter_variables=filter_variables, + ) formatted_test_body = textwrap.indent(test_body, " " if test_framework == "unittest" else " ") @@ -142,3 +196,85 @@ def create_trace_replay_test( test_template += f"def test_{alias}({self}):\n{formatted_test_body}\n" return imports + "\n" + metadata + "\n" + test_template + +def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework: str = "pytest", max_run_count: int = 100) -> None: + """Generate multiple replay tests from the traced function calls, grouping by benchmark name. + + Args: + trace_file_path: Path to the SQLite database file + output_dir: Directory to write the generated tests (if None, only returns the code) + project_root: Root directory of the project for module imports + test_framework: 'pytest' or 'unittest' + max_run_count: Maximum number of runs to include per function + + Returns: + Dictionary mapping benchmark names to generated test code + + """ + try: + # Connect to the database + conn = sqlite3.connect(trace_file_path.as_posix()) + cursor = conn.cursor() + + # Get distinct benchmark names + cursor.execute( + "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM function_calls" + ) + benchmarks = cursor.fetchall() + + # Generate a test for each benchmark + for benchmark in benchmarks: + benchmark_function_name, benchmark_file_name = benchmark + # Get functions associated with this benchmark + cursor.execute( + "SELECT DISTINCT function_name, class_name, module_name, file_name, benchmark_line_number FROM function_calls " + "WHERE benchmark_function_name = ? AND benchmark_file_name = ?", + (benchmark_function_name, benchmark_file_name) + ) + + functions_data = [] + for func_row in cursor.fetchall(): + function_name, class_name, module_name, file_name, benchmark_line_number = func_row + + # Add this function to our list + functions_data.append({ + "function_name": function_name, + "class_name": class_name, + "file_name": file_name, + "module_name": module_name, + "benchmark_function_name": benchmark_function_name, + "benchmark_file_name": benchmark_file_name, + "benchmark_line_number": benchmark_line_number, + "function_properties": inspect_top_level_functions_or_methods( + file_name=file_name, + function_or_method_name=function_name, + class_name=class_name, + ) + }) + + if not functions_data: + print(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") + continue + + # Generate the test code for this benchmark + test_code = create_trace_replay_test_code( + trace_file=trace_file_path.as_posix(), + functions_data=functions_data, + test_framework=test_framework, + max_run_count=max_run_count, + ) + test_code = isort.code(test_code) + + # Write to file if requested + if output_dir: + output_file = get_test_file_path( + test_dir=Path(output_dir), function_name=f"{benchmark_file_name[5:]}_{benchmark_function_name}", test_type="replay" + ) + with open(output_file, 'w') as f: + f.write(test_code) + print(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") + + conn.close() + + except Exception as e: + print(f"Error generating replay tests: {e}") diff --git a/tests/test_instrument_codeflash_capture.py b/tests/test_instrument_codeflash_capture.py index fe5a6bcd3..5cd5ce322 100644 --- a/tests/test_instrument_codeflash_capture.py +++ b/tests/test_instrument_codeflash_capture.py @@ -1,356 +1,193 @@ -from pathlib import Path +from __future__ import annotations -from codeflash.code_utils.code_utils import get_run_tmp_file -from codeflash.discovery.functions_to_optimize import FunctionToOptimize -from codeflash.models.models import FunctionParent -from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture +from codeflash.benchmarking.instrument_codeflash_trace import add_codeflash_decorator_to_code -def test_add_codeflash_capture(): - original_code = """ -class MyClass: - def __init__(self): - self.x = 1 - - def target_function(self): - return self.x + 1 +def test_add_decorator_to_normal_function() -> None: + """Test adding decorator to a normal function.""" + code = """ +def normal_function(): + return "Hello, World!" """ - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - expected = f""" -from codeflash.verification.codeflash_capture import codeflash_capture - - -class MyClass: - - @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) - def __init__(self): - self.x = 1 - def target_function(self): - return self.x + 1 -""" - test_path.write_text(original_code) - - function = FunctionToOptimize( - function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="normal_function" ) - try: - instrument_codeflash_capture(function, {}, test_path.parent) - modified_code = test_path.read_text() - assert modified_code.strip() == expected.strip() - - finally: - test_path.unlink(missing_ok=True) - - -def test_add_codeflash_capture_no_parent(): - original_code = """ -class MyClass: - - def target_function(self): - return self.x + 1 + expected_code = """ +@codeflash_trace +def normal_function(): + return "Hello, World!" """ - expected = """ -class MyClass: + assert modified_code.strip() == expected_code.strip() - def target_function(self): - return self.x + 1 +def test_add_decorator_to_normal_method() -> None: + """Test adding decorator to a normal method.""" + code = """ +class TestClass: + def normal_method(self): + return "Hello from method" """ - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - test_path.write_text(original_code) - - function = FunctionToOptimize(function_name="target_function", file_path=test_path, parents=[]) - try: - instrument_codeflash_capture(function, {}, test_path.parent) - modified_code = test_path.read_text() - assert modified_code.strip() == expected.strip() - finally: - test_path.unlink(missing_ok=True) - - -def test_add_codeflash_capture_no_init(): - # Test input code - original_code = """ -class MyClass(ParentClass): + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="normal_method", + class_name="TestClass" + ) - def target_function(self): - return self.x + 1 + expected_code = """ +class TestClass: + @codeflash_trace + def normal_method(self): + return "Hello from method" """ - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - expected = f""" -from codeflash.verification.codeflash_capture import codeflash_capture + assert modified_code.strip() == expected_code.strip() -class MyClass(ParentClass): - - @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def target_function(self): - return self.x + 1 +def test_add_decorator_to_classmethod() -> None: + """Test adding decorator to a classmethod.""" + code = """ +class TestClass: + @classmethod + def class_method(cls): + return "Hello from classmethod" """ - test_path.write_text(original_code) - function = FunctionToOptimize( - function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="class_method", + class_name="TestClass" ) - try: - instrument_codeflash_capture(function, {}, test_path.parent) - modified_code = test_path.read_text() - assert modified_code.strip() == expected.strip() - - finally: - test_path.unlink(missing_ok=True) - - -def test_add_codeflash_capture_with_helpers(): - # Test input code - original_code = """ -class MyClass: - def __init__(self): - self.x = 1 - - def target_function(self): - return helper() + 1 - - def helper(self): - return self.x + expected_code = """ +class TestClass: + @classmethod + @codeflash_trace + def class_method(cls): + return "Hello from classmethod" """ - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - expected = f""" -from codeflash.verification.codeflash_capture import codeflash_capture + assert modified_code.strip() == expected_code.strip() -class MyClass: - - @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) - def __init__(self): - self.x = 1 - - def target_function(self): - return helper() + 1 - - def helper(self): - return self.x +def test_add_decorator_to_staticmethod() -> None: + """Test adding decorator to a staticmethod.""" + code = """ +class TestClass: + @staticmethod + def static_method(): + return "Hello from staticmethod" """ - test_path.write_text(original_code) - - function = FunctionToOptimize( - function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="static_method", + class_name="TestClass" ) - try: - instrument_codeflash_capture( - function, {test_path: {"MyClass"}}, test_path.parent - ) # MyClass was removed from the file_path_to_helper_class as it shares class with FTO - modified_code = test_path.read_text() - assert modified_code.strip() == expected.strip() - - finally: - test_path.unlink(missing_ok=True) - - -def test_add_codeflash_capture_with_helpers_2(): - # Test input code - original_code = """ -from test_helper_file import HelperClass - -class MyClass: - def __init__(self): - self.x = 1 - - def target_function(self): - return HelperClass().helper() + 1 -""" - original_helper = """ -class HelperClass: - def __init__(self): - self.y = 1 - def helper(self): - return 1 -""" - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - expected = f""" -from test_helper_file import HelperClass - -from codeflash.verification.codeflash_capture import codeflash_capture - - -class MyClass: - - @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) - def __init__(self): - self.x = 1 - - def target_function(self): - return HelperClass().helper() + 1 + expected_code = """ +class TestClass: + @staticmethod + @codeflash_trace + def static_method(): + return "Hello from staticmethod" """ - expected_helper = f""" -from codeflash.verification.codeflash_capture import codeflash_capture - -class HelperClass: + assert modified_code.strip() == expected_code.strip() - @codeflash_capture(function_name='HelperClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) - def __init__(self): - self.y = 1 - - def helper(self): - return 1 +def test_add_decorator_to_init_function() -> None: + """Test adding decorator to an __init__ function.""" + code = """ +class TestClass: + def __init__(self, value): + self.value = value """ - test_path.write_text(original_code) - helper_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_helper_file.py").resolve() - helper_path.write_text(original_helper) - - function = FunctionToOptimize( - function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="__init__", + class_name="TestClass" ) - try: - instrument_codeflash_capture(function, {helper_path: {"HelperClass"}}, test_path.parent) - modified_code = test_path.read_text() - assert modified_code.strip() == expected.strip() - assert helper_path.read_text().strip() == expected_helper.strip() - finally: - test_path.unlink(missing_ok=True) - helper_path.unlink(missing_ok=True) - - -def test_add_codeflash_capture_with_multiple_helpers(): - # Test input code with imports from two helper files - original_code = """ -from helper_file_1 import HelperClass1 -from helper_file_2 import HelperClass2, AnotherHelperClass - -class MyClass: - def __init__(self): - self.x = 1 - - def target_function(self): - helper1 = HelperClass1().helper1() - helper2 = HelperClass2().helper2() - another = AnotherHelperClass().another_helper() - return helper1 + helper2 + another + expected_code = """ +class TestClass: + @codeflash_trace + def __init__(self, value): + self.value = value """ - # First helper file content - original_helper1 = """ -class HelperClass1: - def __init__(self): - self.y = 1 - def helper1(self): - return 1 -""" + assert modified_code.strip() == expected_code.strip() - # Second helper file content - original_helper2 = """ -class HelperClass2: - def __init__(self): - self.z = 2 - def helper2(self): - return 2 - -class AnotherHelperClass: - def another_helper(self): - return 3 +def test_add_decorator_with_multiple_decorators() -> None: + """Test adding decorator to a function with multiple existing decorators.""" + code = """ +class TestClass: + @property + @other_decorator + def property_method(self): + return self._value """ - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - expected = f""" -from helper_file_1 import HelperClass1 -from helper_file_2 import AnotherHelperClass, HelperClass2 -from codeflash.verification.codeflash_capture import codeflash_capture - - -class MyClass: - - @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) - def __init__(self): - self.x = 1 + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="property_method", + class_name="TestClass" + ) - def target_function(self): - helper1 = HelperClass1().helper1() - helper2 = HelperClass2().helper2() - another = AnotherHelperClass().another_helper() - return helper1 + helper2 + another + expected_code = """ +class TestClass: + @property + @other_decorator + @codeflash_trace + def property_method(self): + return self._value """ - # Expected output for first helper file - expected_helper1 = f""" -from codeflash.verification.codeflash_capture import codeflash_capture + assert modified_code.strip() == expected_code.strip() +def test_add_decorator_to_function_in_multiple_classes() -> None: + """Test that only the right class's method gets the decorator.""" + code = """ +class TestClass: + def test_method(self): + return "This should get decorated" -class HelperClass1: - - @codeflash_capture(function_name='HelperClass1.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) - def __init__(self): - self.y = 1 - - def helper1(self): - return 1 +class OtherClass: + def test_method(self): + return "This should NOT get decorated" """ - # Expected output for second helper file - expected_helper2 = f""" -from codeflash.verification.codeflash_capture import codeflash_capture - - -class HelperClass2: - - @codeflash_capture(function_name='HelperClass2.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) - def __init__(self): - self.z = 2 - - def helper2(self): - return 2 - -class AnotherHelperClass: + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="test_method", + class_name="TestClass" + ) - @codeflash_capture(function_name='AnotherHelperClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + expected_code = """ +class TestClass: + @codeflash_trace + def test_method(self): + return "This should get decorated" - def another_helper(self): - return 3 +class OtherClass: + def test_method(self): + return "This should NOT get decorated" """ - # Set up test files - helper1_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/helper_file_1.py").resolve() - helper2_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/helper_file_2.py").resolve() + assert modified_code.strip() == expected_code.strip() - # Write original content to files - test_path.write_text(original_code) - helper1_path.write_text(original_helper1) - helper2_path.write_text(original_helper2) +def test_add_decorator_to_nonexistent_function() -> None: + """Test that code remains unchanged when function doesn't exist.""" + code = """ +def existing_function(): + return "This exists" +""" - # Create FunctionToOptimize instance - function = FunctionToOptimize( - function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="nonexistent_function" ) - try: - # Instrument code with multiple helper files - helper_classes = {helper1_path: {"HelperClass1"}, helper2_path: {"HelperClass2", "AnotherHelperClass"}} - instrument_codeflash_capture(function, helper_classes, test_path.parent) - - # Verify the modifications - modified_code = test_path.read_text() - modified_helper1 = helper1_path.read_text() - modified_helper2 = helper2_path.read_text() - - assert modified_code.strip() == expected.strip() - assert modified_helper1.strip() == expected_helper1.strip() - assert modified_helper2.strip() == expected_helper2.strip() - - finally: - # Clean up test files - test_path.unlink(missing_ok=True) - helper1_path.unlink(missing_ok=True) - helper2_path.unlink(missing_ok=True) + # Code should remain unchanged + assert modified_code.strip() == code.strip() diff --git a/tests/test_instrument_codeflash_trace.py b/tests/test_instrument_codeflash_trace.py new file mode 100644 index 000000000..56008faa9 --- /dev/null +++ b/tests/test_instrument_codeflash_trace.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +from pathlib import Path + +from codeflash.benchmarking.instrument_codeflash_trace import add_codeflash_decorator_to_code + +from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize + + +def test_add_decorator_to_normal_function() -> None: + """Test adding decorator to a normal function.""" + code = """ +def normal_function(): + return "Hello, World!" +""" + + fto = FunctionToOptimize( + function_name="normal_function", + file_path=Path("dummy_path.py"), + parents=[] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +@codeflash_trace +def normal_function(): + return "Hello, World!" +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_normal_method() -> None: + """Test adding decorator to a normal method.""" + code = """ +class TestClass: + def normal_method(self): + return "Hello from method" +""" + + fto = FunctionToOptimize( + function_name="normal_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @codeflash_trace + def normal_method(self): + return "Hello from method" +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_classmethod() -> None: + """Test adding decorator to a classmethod.""" + code = """ +class TestClass: + @classmethod + def class_method(cls): + return "Hello from classmethod" +""" + + fto = FunctionToOptimize( + function_name="class_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @classmethod + @codeflash_trace + def class_method(cls): + return "Hello from classmethod" +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_staticmethod() -> None: + """Test adding decorator to a staticmethod.""" + code = """ +class TestClass: + @staticmethod + def static_method(): + return "Hello from staticmethod" +""" + + fto = FunctionToOptimize( + function_name="static_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @staticmethod + @codeflash_trace + def static_method(): + return "Hello from staticmethod" +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_init_function() -> None: + """Test adding decorator to an __init__ function.""" + code = """ +class TestClass: + def __init__(self, value): + self.value = value +""" + + fto = FunctionToOptimize( + function_name="__init__", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @codeflash_trace + def __init__(self, value): + self.value = value +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_with_multiple_decorators() -> None: + """Test adding decorator to a function with multiple existing decorators.""" + code = """ +class TestClass: + @property + @other_decorator + def property_method(self): + return self._value +""" + + fto = FunctionToOptimize( + function_name="property_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @property + @other_decorator + @codeflash_trace + def property_method(self): + return self._value +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_function_in_multiple_classes() -> None: + """Test that only the right class's method gets the decorator.""" + code = """ +class TestClass: + def test_method(self): + return "This should get decorated" + +class OtherClass: + def test_method(self): + return "This should NOT get decorated" +""" + + fto = FunctionToOptimize( + function_name="test_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @codeflash_trace + def test_method(self): + return "This should get decorated" + +class OtherClass: + def test_method(self): + return "This should NOT get decorated" +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_nonexistent_function() -> None: + """Test that code remains unchanged when function doesn't exist.""" + code = """ +def existing_function(): + return "This exists" +""" + + fto = FunctionToOptimize( + function_name="nonexistent_function", + file_path=Path("dummy_path.py"), + parents=[] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + # Code should remain unchanged + assert modified_code.strip() == code.strip() diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 570888fcc..c49e7c693 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,27 +1,154 @@ +import sqlite3 + from codeflash.benchmarking.codeflash_trace import codeflash_trace from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest +from codeflash.benchmarking.replay_test import generate_replay_test from pathlib import Path from codeflash.code_utils.code_utils import get_run_tmp_file import shutil + 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" - # make directory in project_root / "tests" - - tests_root = project_root / "tests" / "test_trace_benchmarks" tests_root.mkdir(parents=False, exist_ok=False) output_file = (tests_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_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM function_calls ORDER BY benchmark_file_name, benchmark_function_name, function_name") + 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)}" + + # Expected function calls + expected_calls = [ + ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_class_sort", "test_benchmark_bubble_sort.py", 20), + + ("sort_class", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_class_sort", "test_benchmark_bubble_sort.py", 18), + + ("sort_static", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_class_sort", "test_benchmark_bubble_sort.py", 19), + + ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_class_sort", "test_benchmark_bubble_sort.py", 17), + + ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_sort", "test_benchmark_bubble_sort.py", 7), + + ("compute_and_sort", "", "code_to_optimize.process_and_bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/process_and_bubble_sort_codeflash_trace.py'}", + "test_compute_and_sort", "test_process_and_sort.py", 4), + + ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_no_func", "test_process_and_sort.py", 8), + ] + 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_name" + 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_file_name" + assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number" + # Close connection + conn.close() + generate_replay_test(output_file, tests_root) + test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort_py_test_class_sort__replay_test_0.py") + assert test_class_sort_path.exists() + test_class_sort_code = f""" +import dill as pickle + +from code_to_optimize.bubble_sort_codeflash_trace import \\ + Sorter as code_to_optimize_bubble_sort_codeflash_trace_Sorter +from codeflash.benchmarking.replay_test import get_next_arg_and_return + +functions = ['sorter', 'sort_class', 'sort_static'] +trace_file_path = r"{output_file.as_posix()}" + +def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + function_name = "sorter" + if not args: + raise ValueError("No arguments provided for the method.") + if function_name == "__init__": + ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter(*args[1:], **kwargs) + else: + instance = args[0] # self + ret = instance.sorter(*args[1:], **kwargs) + +def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + if not args: + raise ValueError("No arguments provided for the method.") + ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_class(*args[1:], **kwargs) + +def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_static(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_static(*args, **kwargs) + +def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + function_name = "__init__" + if not args: + raise ValueError("No arguments provided for the method.") + if function_name == "__init__": + ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter(*args[1:], **kwargs) + else: + instance = args[0] # self + ret = instance(*args[1:], **kwargs) + +""" + assert test_class_sort_path.read_text("utf-8").strip()==test_class_sort_code.strip() + + test_sort_path = tests_root / Path("test_benchmark_bubble_sort_py_test_sort__replay_test_0.py") + assert test_sort_path.exists() + test_sort_code = f""" +import dill as pickle + +from code_to_optimize.bubble_sort_codeflash_trace import \\ + sorter as code_to_optimize_bubble_sort_codeflash_trace_sorter +from codeflash.benchmarking.replay_test import get_next_arg_and_return + +functions = ['sorter'] +trace_file_path = r"{output_file}" - test1_path = tests_root / Path("test_benchmark_bubble_sort_py_test_sort__replay_test_0.py") - assert test1_path.exists() +def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + ret = code_to_optimize_bubble_sort_codeflash_trace_sorter(*args, **kwargs) - # test1_code = """""" - # assert test1_path.read_text("utf-8").strip()==test1_code.strip() - # cleanup - # shutil.rmtree(tests_root) - # output_file.unlink() \ No newline at end of file +""" + assert test_sort_path.read_text("utf-8").strip()==test_sort_code.strip() + finally: + # cleanup + shutil.rmtree(tests_root) + pass \ No newline at end of file From 4c19e6f7816f902aba6bb24f057f82cafea9a66b Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 14 Mar 2025 15:05:55 -0700 Subject: [PATCH 010/122] restored overwritten logic --- tests/test_instrument_codeflash_capture.py | 441 ++++++++++++++------- 1 file changed, 302 insertions(+), 139 deletions(-) diff --git a/tests/test_instrument_codeflash_capture.py b/tests/test_instrument_codeflash_capture.py index 5cd5ce322..fe5a6bcd3 100644 --- a/tests/test_instrument_codeflash_capture.py +++ b/tests/test_instrument_codeflash_capture.py @@ -1,193 +1,356 @@ -from __future__ import annotations +from pathlib import Path -from codeflash.benchmarking.instrument_codeflash_trace import add_codeflash_decorator_to_code +from codeflash.code_utils.code_utils import get_run_tmp_file +from codeflash.discovery.functions_to_optimize import FunctionToOptimize +from codeflash.models.models import FunctionParent +from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture -def test_add_decorator_to_normal_function() -> None: - """Test adding decorator to a normal function.""" - code = """ -def normal_function(): - return "Hello, World!" +def test_add_codeflash_capture(): + original_code = """ +class MyClass: + def __init__(self): + self.x = 1 + + def target_function(self): + return self.x + 1 """ + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + expected = f""" +from codeflash.verification.codeflash_capture import codeflash_capture - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="normal_function" - ) - expected_code = """ -@codeflash_trace -def normal_function(): - return "Hello, World!" -""" +class MyClass: - assert modified_code.strip() == expected_code.strip() + @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) + def __init__(self): + self.x = 1 -def test_add_decorator_to_normal_method() -> None: - """Test adding decorator to a normal method.""" - code = """ -class TestClass: - def normal_method(self): - return "Hello from method" + def target_function(self): + return self.x + 1 """ + test_path.write_text(original_code) - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="normal_method", - class_name="TestClass" + function = FunctionToOptimize( + function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] ) - expected_code = """ -class TestClass: - @codeflash_trace - def normal_method(self): - return "Hello from method" + try: + instrument_codeflash_capture(function, {}, test_path.parent) + modified_code = test_path.read_text() + assert modified_code.strip() == expected.strip() + + finally: + test_path.unlink(missing_ok=True) + + +def test_add_codeflash_capture_no_parent(): + original_code = """ +class MyClass: + + def target_function(self): + return self.x + 1 """ - assert modified_code.strip() == expected_code.strip() + expected = """ +class MyClass: -def test_add_decorator_to_classmethod() -> None: - """Test adding decorator to a classmethod.""" - code = """ -class TestClass: - @classmethod - def class_method(cls): - return "Hello from classmethod" + def target_function(self): + return self.x + 1 """ + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + test_path.write_text(original_code) - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="class_method", - class_name="TestClass" - ) + function = FunctionToOptimize(function_name="target_function", file_path=test_path, parents=[]) + + try: + instrument_codeflash_capture(function, {}, test_path.parent) + modified_code = test_path.read_text() + assert modified_code.strip() == expected.strip() + finally: + test_path.unlink(missing_ok=True) + + +def test_add_codeflash_capture_no_init(): + # Test input code + original_code = """ +class MyClass(ParentClass): - expected_code = """ -class TestClass: - @classmethod - @codeflash_trace - def class_method(cls): - return "Hello from classmethod" + def target_function(self): + return self.x + 1 """ + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + expected = f""" +from codeflash.verification.codeflash_capture import codeflash_capture - assert modified_code.strip() == expected_code.strip() -def test_add_decorator_to_staticmethod() -> None: - """Test adding decorator to a staticmethod.""" - code = """ -class TestClass: - @staticmethod - def static_method(): - return "Hello from staticmethod" +class MyClass(ParentClass): + + @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def target_function(self): + return self.x + 1 """ + test_path.write_text(original_code) - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="static_method", - class_name="TestClass" + function = FunctionToOptimize( + function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] ) - expected_code = """ -class TestClass: - @staticmethod - @codeflash_trace - def static_method(): - return "Hello from staticmethod" + try: + instrument_codeflash_capture(function, {}, test_path.parent) + modified_code = test_path.read_text() + assert modified_code.strip() == expected.strip() + + finally: + test_path.unlink(missing_ok=True) + + +def test_add_codeflash_capture_with_helpers(): + # Test input code + original_code = """ +class MyClass: + def __init__(self): + self.x = 1 + + def target_function(self): + return helper() + 1 + + def helper(self): + return self.x """ + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + expected = f""" +from codeflash.verification.codeflash_capture import codeflash_capture + + +class MyClass: + + @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) + def __init__(self): + self.x = 1 - assert modified_code.strip() == expected_code.strip() + def target_function(self): + return helper() + 1 -def test_add_decorator_to_init_function() -> None: - """Test adding decorator to an __init__ function.""" - code = """ -class TestClass: - def __init__(self, value): - self.value = value + def helper(self): + return self.x """ - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="__init__", - class_name="TestClass" + test_path.write_text(original_code) + + function = FunctionToOptimize( + function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] ) - expected_code = """ -class TestClass: - @codeflash_trace - def __init__(self, value): - self.value = value + try: + instrument_codeflash_capture( + function, {test_path: {"MyClass"}}, test_path.parent + ) # MyClass was removed from the file_path_to_helper_class as it shares class with FTO + modified_code = test_path.read_text() + assert modified_code.strip() == expected.strip() + + finally: + test_path.unlink(missing_ok=True) + + +def test_add_codeflash_capture_with_helpers_2(): + # Test input code + original_code = """ +from test_helper_file import HelperClass + +class MyClass: + def __init__(self): + self.x = 1 + + def target_function(self): + return HelperClass().helper() + 1 """ + original_helper = """ +class HelperClass: + def __init__(self): + self.y = 1 + def helper(self): + return 1 +""" + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + expected = f""" +from test_helper_file import HelperClass + +from codeflash.verification.codeflash_capture import codeflash_capture + + +class MyClass: + + @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) + def __init__(self): + self.x = 1 + + def target_function(self): + return HelperClass().helper() + 1 +""" + expected_helper = f""" +from codeflash.verification.codeflash_capture import codeflash_capture - assert modified_code.strip() == expected_code.strip() -def test_add_decorator_with_multiple_decorators() -> None: - """Test adding decorator to a function with multiple existing decorators.""" - code = """ -class TestClass: - @property - @other_decorator - def property_method(self): - return self._value +class HelperClass: + + @codeflash_capture(function_name='HelperClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) + def __init__(self): + self.y = 1 + + def helper(self): + return 1 """ - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="property_method", - class_name="TestClass" + test_path.write_text(original_code) + helper_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_helper_file.py").resolve() + helper_path.write_text(original_helper) + + function = FunctionToOptimize( + function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] ) - expected_code = """ -class TestClass: - @property - @other_decorator - @codeflash_trace - def property_method(self): - return self._value + try: + instrument_codeflash_capture(function, {helper_path: {"HelperClass"}}, test_path.parent) + modified_code = test_path.read_text() + assert modified_code.strip() == expected.strip() + assert helper_path.read_text().strip() == expected_helper.strip() + finally: + test_path.unlink(missing_ok=True) + helper_path.unlink(missing_ok=True) + + +def test_add_codeflash_capture_with_multiple_helpers(): + # Test input code with imports from two helper files + original_code = """ +from helper_file_1 import HelperClass1 +from helper_file_2 import HelperClass2, AnotherHelperClass + +class MyClass: + def __init__(self): + self.x = 1 + + def target_function(self): + helper1 = HelperClass1().helper1() + helper2 = HelperClass2().helper2() + another = AnotherHelperClass().another_helper() + return helper1 + helper2 + another +""" + + # First helper file content + original_helper1 = """ +class HelperClass1: + def __init__(self): + self.y = 1 + def helper1(self): + return 1 +""" + + # Second helper file content + original_helper2 = """ +class HelperClass2: + def __init__(self): + self.z = 2 + def helper2(self): + return 2 + +class AnotherHelperClass: + def another_helper(self): + return 3 """ + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + expected = f""" +from helper_file_1 import HelperClass1 +from helper_file_2 import AnotherHelperClass, HelperClass2 - assert modified_code.strip() == expected_code.strip() +from codeflash.verification.codeflash_capture import codeflash_capture -def test_add_decorator_to_function_in_multiple_classes() -> None: - """Test that only the right class's method gets the decorator.""" - code = """ -class TestClass: - def test_method(self): - return "This should get decorated" -class OtherClass: - def test_method(self): - return "This should NOT get decorated" +class MyClass: + + @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) + def __init__(self): + self.x = 1 + + def target_function(self): + helper1 = HelperClass1().helper1() + helper2 = HelperClass2().helper2() + another = AnotherHelperClass().another_helper() + return helper1 + helper2 + another """ - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="test_method", - class_name="TestClass" - ) + # Expected output for first helper file + expected_helper1 = f""" +from codeflash.verification.codeflash_capture import codeflash_capture + - expected_code = """ -class TestClass: - @codeflash_trace - def test_method(self): - return "This should get decorated" +class HelperClass1: -class OtherClass: - def test_method(self): - return "This should NOT get decorated" + @codeflash_capture(function_name='HelperClass1.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) + def __init__(self): + self.y = 1 + + def helper1(self): + return 1 """ - assert modified_code.strip() == expected_code.strip() + # Expected output for second helper file + expected_helper2 = f""" +from codeflash.verification.codeflash_capture import codeflash_capture + + +class HelperClass2: -def test_add_decorator_to_nonexistent_function() -> None: - """Test that code remains unchanged when function doesn't exist.""" - code = """ -def existing_function(): - return "This exists" + @codeflash_capture(function_name='HelperClass2.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) + def __init__(self): + self.z = 2 + + def helper2(self): + return 2 + +class AnotherHelperClass: + + @codeflash_capture(function_name='AnotherHelperClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def another_helper(self): + return 3 """ - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="nonexistent_function" + # Set up test files + helper1_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/helper_file_1.py").resolve() + helper2_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/helper_file_2.py").resolve() + + # Write original content to files + test_path.write_text(original_code) + helper1_path.write_text(original_helper1) + helper2_path.write_text(original_helper2) + + # Create FunctionToOptimize instance + function = FunctionToOptimize( + function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] ) - # Code should remain unchanged - assert modified_code.strip() == code.strip() + try: + # Instrument code with multiple helper files + helper_classes = {helper1_path: {"HelperClass1"}, helper2_path: {"HelperClass2", "AnotherHelperClass"}} + instrument_codeflash_capture(function, helper_classes, test_path.parent) + + # Verify the modifications + modified_code = test_path.read_text() + modified_helper1 = helper1_path.read_text() + modified_helper2 = helper2_path.read_text() + + assert modified_code.strip() == expected.strip() + assert modified_helper1.strip() == expected_helper1.strip() + assert modified_helper2.strip() == expected_helper2.strip() + + finally: + # Clean up test files + test_path.unlink(missing_ok=True) + helper1_path.unlink(missing_ok=True) + helper2_path.unlink(missing_ok=True) From 7eba0317c1ea8105a8bba7b79110b4e15a645a59 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 18 Mar 2025 10:00:05 -0700 Subject: [PATCH 011/122] functioning end to end, gets the funciton impact on benchmarks --- .../benchmarking/benchmark_database_utils.py | 179 ++++++++++++++++++ codeflash/benchmarking/codeflash_trace.py | 124 +----------- codeflash/benchmarking/get_trace_info.py | 165 ++++++++-------- .../instrument_codeflash_trace.py | 37 +++- codeflash/benchmarking/plugin/plugin.py | 20 +- .../pytest_new_process_trace_benchmarks.py | 16 +- codeflash/benchmarking/replay_test.py | 5 +- codeflash/benchmarking/trace_benchmarks.py | 6 +- codeflash/discovery/functions_to_optimize.py | 43 +++-- .../pytest_new_process_discover_benchmarks.py | 54 ------ codeflash/optimization/function_optimizer.py | 17 +- codeflash/optimization/optimizer.py | 25 ++- tests/test_instrument_codeflash_trace.py | 7 + 13 files changed, 394 insertions(+), 304 deletions(-) create mode 100644 codeflash/benchmarking/benchmark_database_utils.py delete mode 100644 codeflash/discovery/pytest_new_process_discover_benchmarks.py diff --git a/codeflash/benchmarking/benchmark_database_utils.py b/codeflash/benchmarking/benchmark_database_utils.py new file mode 100644 index 000000000..b9b36079d --- /dev/null +++ b/codeflash/benchmarking/benchmark_database_utils.py @@ -0,0 +1,179 @@ +import sqlite3 +from pathlib import Path + +import pickle + + +class BenchmarkDatabaseUtils: + def __init__(self, trace_path :Path) -> None: + self.trace_path = trace_path + self.connection = None + + def setup(self) -> None: + try: + # Open connection + self.connection = sqlite3.connect(self.trace_path) + cur = self.connection.cursor() + cur.execute("PRAGMA synchronous = OFF") + cur.execute( + "CREATE TABLE IF NOT EXISTS function_calls(" + "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," + "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," + "time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" + ) + cur.execute( + "CREATE TABLE IF NOT EXISTS benchmark_timings(" + "benchmark_file_name TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER," + "time_ns INTEGER)" # Added closing parenthesis + ) + self.connection.commit() + # Don't close the connection here + except Exception as e: + print(f"Database setup error: {e}") + if self.connection: + self.connection.close() + self.connection = None + raise + + def write_function_timings(self, data: list[tuple]) -> None: + if not self.connection: + self.connection = sqlite3.connect(self.trace_path) + + try: + cur = self.connection.cursor() + # Insert data into the function_calls table + cur.executemany( + "INSERT INTO function_calls " + "(function_name, class_name, module_name, file_name, benchmark_function_name, " + "benchmark_file_name, benchmark_line_number, time_ns, overhead_time_ns, args, kwargs) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + data + ) + self.connection.commit() + except Exception as e: + print(f"Error writing to function timings database: {e}") + self.connection.rollback() + raise + + def write_benchmark_timings(self, data: list[tuple]) -> None: + if not self.connection: + self.connection = sqlite3.connect(self.trace_path) + + try: + cur = self.connection.cursor() + # Insert data into the benchmark_timings table + cur.executemany( + "INSERT INTO benchmark_timings (benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns) VALUES (?, ?, ?, ?)", + data + ) + self.connection.commit() + except Exception as e: + print(f"Error writing to benchmark timings database: {e}") + self.connection.rollback() + raise + + def print_function_timings(self, limit: int = None) -> None: + """Print the contents of a CodeflashTrace SQLite database. + + Args: + limit: Maximum number of records to print (None for all) + """ + if not self.connection: + self.connection = sqlite3.connect(self.trace_path) + try: + cur = self.connection.cursor() + + # Get the count of records + cur.execute("SELECT COUNT(*) FROM function_calls") + total_records = cur.fetchone()[0] + print(f"Found {total_records} function call records in {self.trace_path}") + + # Build the query with optional limit + query = "SELECT * FROM function_calls" + if limit: + query += f" LIMIT {limit}" + + # Execute the query + cur.execute(query) + + # Print column names + columns = [desc[0] for desc in cur.description] + print("\nColumns:", columns) + print("\n" + "=" * 80 + "\n") + + # Print each row + for i, row in enumerate(cur.fetchall()): + print(f"Record #{i + 1}:") + print(f" Function: {row[0]}") + print(f" Class: {row[1]}") + print(f" Module: {row[2]}") + print(f" File: {row[3]}") + print(f" Benchmark Function: {row[4] or 'N/A'}") + print(f" Benchmark File: {row[5] or 'N/A'}") + print(f" Benchmark Line: {row[6] or 'N/A'}") + print(f" Execution Time: {row[7]:.6f} seconds") + print(f" Overhead Time: {row[8]:.6f} seconds") + + # Unpickle and print args and kwargs + try: + args = pickle.loads(row[9]) + kwargs = pickle.loads(row[10]) + + print(f" Args: {args}") + print(f" Kwargs: {kwargs}") + except Exception as e: + print(f" Error unpickling args/kwargs: {e}") + print(f" Raw args: {row[9]}") + print(f" Raw kwargs: {row[10]}") + + print("\n" + "-" * 40 + "\n") + + except Exception as e: + print(f"Error reading database: {e}") + + def print_benchmark_timings(self, limit: int = None) -> None: + """Print the contents of a CodeflashTrace SQLite database. + Args: + limit: Maximum number of records to print (None for all) + """ + if not self.connection: + self.connection = sqlite3.connect(self.trace_path) + try: + cur = self.connection.cursor() + + # Get the count of records + cur.execute("SELECT COUNT(*) FROM benchmark_timings") + total_records = cur.fetchone()[0] + print(f"Found {total_records} benchmark timing records in {self.trace_path}") + + # Build the query with optional limit + query = "SELECT * FROM benchmark_timings" + if limit: + query += f" LIMIT {limit}" + + # Execute the query + cur.execute(query) + + # Print column names + columns = [desc[0] for desc in cur.description] + print("\nColumns:", columns) + print("\n" + "=" * 80 + "\n") + + # Print each row + for i, row in enumerate(cur.fetchall()): + print(f"Record #{i + 1}:") + print(f" Benchmark File: {row[0] or 'N/A'}") + print(f" Benchmark Function: {row[1] or 'N/A'}") + print(f" Benchmark Line: {row[2] or 'N/A'}") + print(f" Execution Time: {row[3] / 1e9:.6f} seconds") # Convert nanoseconds to seconds + print("\n" + "-" * 40 + "\n") + + except Exception as e: + print(f"Error reading benchmark timings database: {e}") + + + def close(self) -> None: + if self.connection: + self.connection.close() + self.connection = None + diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 65ba98783..9b9afead7 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -39,16 +39,15 @@ def __call__(self, func: Callable) -> Callable: @functools.wraps(func) def wrapper(*args, **kwargs): # Measure execution time - start_time = time.time() + start_time = time.perf_counter_ns() result = func(*args, **kwargs) - end_time = time.time() + end_time = time.perf_counter_ns() # Calculate execution time execution_time = end_time - start_time # Measure overhead - overhead_start_time = time.time() - overhead_time = 0 + overhead_start_time = time.perf_counter_ns() try: # Check if currently in pytest benchmark fixture @@ -63,15 +62,16 @@ def wrapper(*args, **kwargs): benchmark_function_name = os.environ.get("CODEFLASH_BENCHMARK_FUNCTION_NAME", "") benchmark_file_name = os.environ.get("CODEFLASH_BENCHMARK_FILE_NAME", "") benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "") - - # Calculate overhead time - overhead_end_time = time.time() - overhead_time = overhead_end_time - overhead_start_time - + # Get class name class_name = "" qualname = func.__qualname__ if "." in qualname: class_name = qualname.split(".")[0] + # Calculate overhead time + overhead_end_time = time.perf_counter_ns() + overhead_time = overhead_end_time - overhead_start_time + + self.function_calls_data.append( (func.__name__, class_name, func.__module__, func.__code__.co_filename, benchmark_function_name, benchmark_file_name, benchmark_line_number, execution_time, @@ -84,114 +84,8 @@ def wrapper(*args, **kwargs): return result return wrapper - def write_to_db(self, output_file: str) -> None: - """Write all collected function call data to the SQLite database. - Args: - output_file: Path to the SQLite database file where results will be stored - """ - if not self.function_calls_data: - print("No function call data to write") - return - self.db_path = output_file - 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, module_name TEXT, file_name TEXT," - "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," - "time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" - ) - - # Insert all data at once - cur.executemany( - "INSERT INTO function_calls " - "(function_name, class_name, module_name, file_name, benchmark_function_name, " - "benchmark_file_name, benchmark_line_number, time_ns, overhead_time_ns, args, kwargs) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - self.function_calls_data - ) - - # Commit and close - con.commit() - con.close() - - print(f"Successfully wrote {len(self.function_calls_data)} function call records to {output_file}") - - # Clear the data after writing - self.function_calls_data.clear() - - except Exception as e: - print(f"Error writing function calls to database: {e}") - - def print_codeflash_db(self, limit: int = None) -> None: - """ - Print the contents of a CodeflashTrace SQLite database. - - Args: - db_path: Path to the SQLite database file - limit: Maximum number of records to print (None for all) - """ - try: - # Connect to the database - conn = sqlite3.connect(self.db_path) - cursor = conn.cursor() - - # Get the count of records - cursor.execute("SELECT COUNT(*) FROM function_calls") - total_records = cursor.fetchone()[0] - print(f"Found {total_records} function call records in {self.db_path}") - - # Build the query with optional limit - query = "SELECT * FROM function_calls" - if limit: - query += f" LIMIT {limit}" - - # Execute the query - cursor.execute(query) - - # Print column names - columns = [desc[0] for desc in cursor.description] - print("\nColumns:", columns) - print("\n" + "=" * 80 + "\n") - - # Print each row - for i, row in enumerate(cursor.fetchall()): - print(f"Record #{i + 1}:") - print(f" Function: {row[0]}") - print(f" Class: {row[1]}") - print(f" Module: {row[2]}") - print(f" File: {row[3]}") - print(f" Benchmark Function: {row[4] or 'N/A'}") - print(f" Benchmark File: {row[5] or 'N/A'}") - print(f" Benchmark Line: {row[6] or 'N/A'}") - print(f" Execution Time: {row[7]:.6f} seconds") - print(f" Overhead Time: {row[8]:.6f} seconds") - - # Unpickle and print args and kwargs - try: - args = pickle.loads(row[9]) - kwargs = pickle.loads(row[10]) - - print(f" Args: {args}") - print(f" Kwargs: {kwargs}") - except Exception as e: - print(f" Error unpickling args/kwargs: {e}") - print(f" Raw args: {row[8]}") - print(f" Raw kwargs: {row[9]}") - - print("\n" + "-" * 40 + "\n") - - conn.close() - - except Exception as e: - print(f"Error reading database: {e}") # Create a singleton instance diff --git a/codeflash/benchmarking/get_trace_info.py b/codeflash/benchmarking/get_trace_info.py index 3dd3831ce..e9a050b84 100644 --- a/codeflash/benchmarking/get_trace_info.py +++ b/codeflash/benchmarking/get_trace_info.py @@ -1,114 +1,109 @@ import sqlite3 from pathlib import Path -from typing import Dict, Set from codeflash.discovery.functions_to_optimize import FunctionToOptimize -def get_function_benchmark_timings(trace_dir: Path, all_functions_to_optimize: list[FunctionToOptimize]) -> dict[str, dict[str, float]]: - """Process all trace files in the given directory and extract timing data for the specified functions. +def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, float]]: + """Process the trace file and extract timing data for all functions. Args: - trace_dir: Path to the directory containing .trace files - all_functions_to_optimize: Set of FunctionToOptimize objects representing functions to include + trace_path: Path to the trace file + all_functions_to_optimize: List of FunctionToOptimize objects (not used directly, + but kept for backward compatibility) Returns: A nested dictionary where: - - Outer keys are function qualified names with file name - - Inner keys are benchmark names (trace filename without .trace extension) + - Outer keys are module_name.qualified_name (module.class.function) + - Inner keys are benchmark filename :: benchmark test function :: line number - Values are function timing in milliseconds """ - # Create a mapping of (filename, function_name, class_name) -> qualified_name for efficient lookups - function_lookup = {} - function_benchmark_timings = {} + # Initialize the result dictionary + result = {} + + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + + try: + # Query the function_calls table for all function calls + cursor.execute( + "SELECT module_name, class_name, function_name, " + "benchmark_file_name, benchmark_function_name, benchmark_line_number, " + "(time_ns - overhead_time_ns) as actual_time_ns " + "FROM function_calls" + ) + + # Process each row + for row in cursor.fetchall(): + module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row + + # Create the function key (module_name.class_name.function_name) + if class_name: + qualified_name = f"{module_name}.{class_name}.{function_name}" + else: + qualified_name = f"{module_name}.{function_name}" - for func in all_functions_to_optimize: - qualified_name = func.qualified_name_with_file_name + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - # Extract components (assumes Path.name gives only filename without directory) - filename = func.file_path - function_name = func.function_name + # Initialize the inner dictionary if needed + if qualified_name not in result: + result[qualified_name] = {} - # Get class name if there's a parent - class_name = func.parents[0].name if func.parents else None + # If multiple calls to the same function in the same benchmark, + # add the times together + if benchmark_key in result[qualified_name]: + result[qualified_name][benchmark_key] += time_ns + else: + result[qualified_name][benchmark_key] = time_ns - # Store in lookup dictionary - key = (filename, function_name, class_name) - function_lookup[key] = qualified_name - function_benchmark_timings[qualified_name] = {} + finally: + # Close the connection + connection.close() - # Find all .trace files in the directory - trace_files = list(trace_dir.glob("*.trace")) + return result - for trace_file in trace_files: - # Extract benchmark name from filename (without .trace) - benchmark_name = trace_file.stem - # Connect to the trace database - conn = sqlite3.connect(trace_file) - cursor = conn.cursor() +def get_benchmark_timings(trace_path: Path) -> dict[str, float]: + """Extract total benchmark timings from trace files. - # For each function we're interested in, query the database directly - for (filename, function_name, class_name), qualified_name in function_lookup.items(): - # Adjust query based on whether we have a class name - if class_name: - cursor.execute( - "SELECT cumulative_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND class_name = ?", - (f"%{filename}", function_name, class_name) - ) - else: - cursor.execute( - "SELECT cumulative_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND (class_name IS NULL OR class_name = '')", - (f"%{filename}", function_name) - ) + Args: + trace_path: Path to the trace file - result = cursor.fetchall() - if len(result) > 1: - print(f"Multiple results found for {qualified_name} in {benchmark_name}: {result}") - if result: - time_ns = result[0][0] - function_benchmark_timings[qualified_name][benchmark_name] = time_ns / 1e6 # Convert to milliseconds + Returns: + A dictionary mapping where: + - Keys are benchmark filename :: benchmark test function :: line number + - Values are total benchmark timing in milliseconds - conn.close() + """ + # Initialize the result dictionary + result = {} - return function_benchmark_timings + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + try: + # Query the benchmark_timings table + cursor.execute( + "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns " + "FROM benchmark_timings" + ) -def get_benchmark_timings(trace_dir: Path) -> dict[str, float]: - """Extract total benchmark timings from trace files. + # Process each row + for row in cursor.fetchall(): + benchmark_file, benchmark_func, benchmark_line, time_ns = row - Args: - trace_dir: Path to the directory containing .trace files + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - Returns: - A dictionary mapping benchmark names to their total execution time in milliseconds. - """ - benchmark_timings = {} - - # Find all .trace files in the directory - trace_files = list(trace_dir.glob("*.trace")) - - for trace_file in trace_files: - # Extract benchmark name from filename (without .trace extension) - benchmark_name = trace_file.stem - - # Connect to the trace database - conn = sqlite3.connect(trace_file) - cursor = conn.cursor() - - # Query the total_time table for the benchmark's total execution time - try: - cursor.execute("SELECT time_ns FROM total_time") - result = cursor.fetchone() - if result: - time_ns = result[0] - # Convert nanoseconds to milliseconds - benchmark_timings[benchmark_name] = time_ns / 1e6 - except sqlite3.OperationalError: - # Handle case where total_time table might not exist - print(f"Warning: Could not get total time for benchmark {benchmark_name}") - - conn.close() - - return benchmark_timings + # Store the timing + result[benchmark_key] = time_ns + + finally: + # Close the connection + connection.close() + + return result diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py index 99b2dad20..93f51baed 100644 --- a/codeflash/benchmarking/instrument_codeflash_trace.py +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -1,3 +1,4 @@ +import isort import libcst as cst from codeflash.discovery.functions_to_optimize import FunctionToOptimize @@ -9,6 +10,7 @@ def __init__(self, function_name, class_name=None): self.function_name = function_name self.class_name = class_name self.in_target_class = (class_name is None) # If no class name, always "in target class" + self.added_codeflash_trace = False def leave_ClassDef(self, original_node, updated_node): if self.class_name and original_node.name.value == self.class_name: @@ -31,12 +33,39 @@ def leave_FunctionDef(self, original_node, updated_node): # Add the new decorator after any existing decorators updated_decorators = list(updated_node.decorators) + [decorator] - + self.added_codeflash_trace = True # Return the updated node with the new decorator return updated_node.with_changes( decorators=updated_decorators ) + def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: + # Create import statement for codeflash_trace + if not self.added_codeflash_trace: + return updated_node + import_stmt = cst.SimpleStatementLine( + body=[ + cst.ImportFrom( + module=cst.Attribute( + value=cst.Attribute( + value=cst.Name(value="codeflash"), + attr=cst.Name(value="benchmarking") + ), + attr=cst.Name(value="codeflash_trace") + ), + names=[ + cst.ImportAlias( + name=cst.Name(value="codeflash_trace") + ) + ] + ) + ] + ) + + # Insert at the beginning of the file + new_body = [import_stmt, *list(updated_node.body)] + + return updated_node.with_changes(body=new_body) def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToOptimize) -> str: """Add codeflash_trace to a function. @@ -63,7 +92,7 @@ def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToO return modified_module.code -def instrument_codeflash_trace( +def instrument_codeflash_trace_decorator( function_to_optimize: FunctionToOptimize ) -> None: """Instrument __init__ function with codeflash_trace decorator if it's in a class.""" @@ -71,10 +100,10 @@ def instrument_codeflash_trace( original_code = function_to_optimize.file_path.read_text(encoding="utf-8") # Modify the code - modified_code = add_codeflash_decorator_to_code( + modified_code = isort.code(add_codeflash_decorator_to_code( original_code, function_to_optimize - ) + )) # Write the modified code back to the file function_to_optimize.file_path.write_text(modified_code, encoding="utf-8") diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index caf175a4e..a5f82fc3a 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -3,7 +3,8 @@ import pytest import time import os -class CodeFlashPlugin: +class CodeFlashBenchmarkPlugin: + benchmark_timings = [] @staticmethod def pytest_addoption(parser): parser.addoption( @@ -38,15 +39,20 @@ def benchmark(request): class Benchmark: def __call__(self, func, *args, **kwargs): - os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = request.node.name - os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = request.node.fspath.basename - os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack + benchmark_file_name = request.node.fspath.basename + benchmark_function_name = request.node.name + line_number = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack + os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name + os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = benchmark_file_name + os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = line_number os.environ["CODEFLASH_BENCHMARKING"] = "True" - start = time.process_time_ns() + + start = time.perf_counter_ns() result = func(*args, **kwargs) - end = time.process_time_ns() + end = time.perf_counter_ns() + os.environ["CODEFLASH_BENCHMARKING"] = "False" - print(f"Benchmark: {func.__name__} took {end - start} ns") + CodeFlashBenchmarkPlugin.benchmark_timings.append((benchmark_file_name, benchmark_function_name, line_number, end - start)) return result return Benchmark() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 04c5e67ea..a83196758 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -1,25 +1,31 @@ import sys from pathlib import Path +from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils from codeflash.verification.verification_utils import get_test_file_path -from plugin.plugin import CodeFlashPlugin +from plugin.plugin import CodeFlashBenchmarkPlugin from codeflash.benchmarking.codeflash_trace import codeflash_trace from codeflash.code_utils.code_utils import get_run_tmp_file benchmarks_root = sys.argv[1] tests_root = sys.argv[2] -output_file = sys.argv[3] +trace_file = sys.argv[3] # current working directory project_root = Path.cwd() if __name__ == "__main__": import pytest try: + db = BenchmarkDatabaseUtils(trace_path=Path(trace_file)) + db.setup() exitcode = pytest.main( - [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashPlugin()] + [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashBenchmarkPlugin()] ) - codeflash_trace.write_to_db(output_file) - codeflash_trace.print_codeflash_db() + db.write_function_timings(codeflash_trace.function_calls_data) + db.write_benchmark_timings(CodeFlashBenchmarkPlugin.benchmark_timings) + db.print_function_timings() + db.print_benchmark_timings() + db.close() except Exception as e: print(f"Failed to collect tests: {e!s}") diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 58ce456c2..75ef7e96d 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -270,8 +270,9 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework output_file = get_test_file_path( test_dir=Path(output_dir), function_name=f"{benchmark_file_name[5:]}_{benchmark_function_name}", test_type="replay" ) - with open(output_file, 'w') as f: - f.write(test_code) + # Write test code to file, parents = true + output_dir.mkdir(parents=True, exist_ok=True) + output_file.write_text(test_code, "utf-8") print(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") conn.close() diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 5c0a077dc..9ae69495d 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -3,19 +3,21 @@ from pathlib import Path import subprocess -def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, output_file: Path) -> None: +def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, trace_file: Path) -> None: + # set up .trace databases result = subprocess.run( [ SAFE_SYS_EXECUTABLE, Path(__file__).parent / "pytest_new_process_trace_benchmarks.py", benchmarks_root, tests_root, - output_file, + trace_file, ], cwd=project_root, check=False, capture_output=True, text=True, + env={"PYTHONPATH": str(project_root)}, ) print("stdout:", result.stdout) print("stderr:", result.stderr) diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index 774571de3..cd1e53f9b 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -143,11 +143,11 @@ def qualified_name(self) -> str: def qualified_name_with_modules_from_root(self, project_root_path: Path) -> str: return f"{module_name_from_file_path(self.file_path, project_root_path)}.{self.qualified_name}" - - @property - def qualified_name_with_file_name(self) -> str: - class_name = self.parents[0].name if self.parents else None - return f"{self.file_path}:{(class_name + ':' if class_name else '')}{self.function_name}" + # + # @property + # def qualified_name_with_file_name(self) -> str: + # class_name = self.parents[0].name if self.parents else None + # return f"{self.file_path}:{(class_name + ':' if class_name else '')}{self.function_name}" def get_functions_to_optimize( @@ -363,23 +363,28 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None: for decorator in body_node.decorator_list ): self.is_classmethod = True - return - else: - # search if the class has a staticmethod with the same name and on the same line number - for body_node in node.body: - if ( - isinstance(body_node, ast.FunctionDef) - and body_node.name == self.function_name - and body_node.lineno in {self.line_no, self.line_no + 1} - and any( + elif any( isinstance(decorator, ast.Name) and decorator.id == "staticmethod" for decorator in body_node.decorator_list - ) - ): - self.is_staticmethod = True - self.is_top_level = True - self.class_name = node.name + ): + self.is_staticmethod = True return + # else: + # # search if the class has a staticmethod with the same name and on the same line number + # for body_node in node.body: + # if ( + # isinstance(body_node, ast.FunctionDef) + # and body_node.name == self.function_name + # # and body_node.lineno in {self.line_no, self.line_no + 1} + # and any( + # isinstance(decorator, ast.Name) and decorator.id == "staticmethod" + # for decorator in body_node.decorator_list + # ) + # ): + # self.is_staticmethod = True + # self.is_top_level = True + # self.class_name = node.name + # return return diff --git a/codeflash/discovery/pytest_new_process_discover_benchmarks.py b/codeflash/discovery/pytest_new_process_discover_benchmarks.py deleted file mode 100644 index 83175218b..000000000 --- a/codeflash/discovery/pytest_new_process_discover_benchmarks.py +++ /dev/null @@ -1,54 +0,0 @@ -import sys -from typing import Any - -# This script should not have any relation to the codeflash package, be careful with imports -cwd = sys.argv[1] -tests_root = sys.argv[2] -pickle_path = sys.argv[3] -collected_tests = [] -pytest_rootdir = None -sys.path.insert(1, str(cwd)) - - -class PytestCollectionPlugin: - def pytest_collection_finish(self, session) -> None: - global pytest_rootdir - collected_tests.extend(session.items) - pytest_rootdir = session.config.rootdir - - -def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, str]]: - test_results = [] - for test in pytest_tests: - test_class = None - if test.cls: - test_class = test.parent.name - - # Determine if this is a benchmark test by checking for the benchmark fixture - is_benchmark = hasattr(test, 'fixturenames') and 'benchmark' in test.fixturenames - test_type = 'benchmark' if is_benchmark else 'regular' - - test_results.append({ - "test_file": str(test.path), - "test_class": test_class, - "test_function": test.name, - "test_type": test_type - }) - return test_results - - -if __name__ == "__main__": - import pytest - - try: - exitcode = pytest.main( - [tests_root, "-pno:logging", "--collect-only", "-m", "not skip"], plugins=[PytestCollectionPlugin()] - ) - except Exception as e: - print(f"Failed to collect tests: {e!s}") - exitcode = -1 - tests = parse_pytest_collection_results(collected_tests) - import pickle - - with open(pickle_path, "wb") as f: - pickle.dump((exitcode, tests, pytest_rootdir), f, protocol=pickle.HIGHEST_PROTOCOL) diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index fd03ab853..b8814b434 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -275,16 +275,23 @@ def optimize_function(self) -> Result[BestOptimization, str]: speedup = explanation.speedup # if self.args.benchmark: original_replay_timing = original_code_baseline.benchmarking_test_results.total_replay_test_runtime() - fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_file_name] - for benchmark_name, og_benchmark_timing in fto_benchmark_timings.items(): - print(f"Calculating speedup for benchmark {benchmark_name}") - total_benchmark_timing = self.total_benchmark_timings[benchmark_name] + fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)] + for benchmark_key, og_benchmark_timing in fto_benchmark_timings.items(): + # benchmark key is benchmark filename :: benchmark test function :: line number + try: + benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") + except ValueError: + print(f"Benchmark key {benchmark_key} is not in the expected format.") + continue + print(f"Calculating speedup for benchmark {benchmark_key}") + total_benchmark_timing = self.total_benchmark_timings[benchmark_key] # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values + print(f"Original benchmark timing: {total_benchmark_timing}") replay_speedup = original_replay_timing / best_optimization.replay_runtime - 1 print(f"Replay speedup: {replay_speedup}") expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / (replay_speedup + 1) * og_benchmark_timing print(f"Expected new benchmark timing: {expected_new_benchmark_timing}") - print(f"Original benchmark timing: {total_benchmark_timing}") + benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 print(f"Benchmark speedup: {benchmark_speedup_percent:.2f}%") diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 26fc70aa7..953f0b755 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient +from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.utils import print_benchmark_table from codeflash.cli_cmds.console import console, logger @@ -25,7 +26,8 @@ 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 codeflash.benchmarking.codeflash_trace import codeflash_trace +from codeflash.benchmarking.instrument_codeflash_trace import instrument_codeflash_trace_decorator + from collections import defaultdict @@ -94,6 +96,8 @@ def run(self) -> None: project_root=self.args.project_root, module_root=self.args.module_root, ) + all_functions_to_optimize = [ + fto for functions_to_optimize in file_to_funcs_to_optimize.values() for fto in functions_to_optimize] if self.args.benchmark: # Insert decorator file_path_to_source_code = defaultdict(str) @@ -103,9 +107,14 @@ def run(self) -> None: 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 + instrument_codeflash_trace_decorator(fto) + trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" + trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Simply run all tests that use pytest-benchmark + generate_replay_test(trace_file, Path(self.args.tests_root) / "codeflash_replay_tests" ) + function_benchmark_timings = get_function_benchmark_timings(trace_file) + total_benchmark_timings = get_benchmark_timings(trace_file) + print(function_benchmark_timings) + print(total_benchmark_timings) logger.info("Finished tracing existing benchmarks") except Exception as e: logger.info(f"Error while tracing existing benchmarks: {e}") @@ -116,13 +125,13 @@ def run(self) -> None: with file.open("w", encoding="utf8") as f: f.write(file_path_to_source_code[file]) - codeflash_trace.print_trace_info() + # 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) - + # return optimizations_found: int = 0 function_iterator_count: int = 0 if self.args.test_framework == "pytest": @@ -207,6 +216,10 @@ def run(self) -> None: function_optimizer = self.create_function_optimizer( function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code, function_benchmark_timings, total_benchmark_timings ) + # function_optimizer = self.create_function_optimizer( + # function_to_optimize, function_to_optimize_ast, function_to_tests, + # validated_original_code[original_module_path].source_code + # ) else: function_optimizer = self.create_function_optimizer( function_to_optimize, function_to_optimize_ast, function_to_tests, diff --git a/tests/test_instrument_codeflash_trace.py b/tests/test_instrument_codeflash_trace.py index 56008faa9..967d5d6f0 100644 --- a/tests/test_instrument_codeflash_trace.py +++ b/tests/test_instrument_codeflash_trace.py @@ -26,6 +26,7 @@ def normal_function(): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace @codeflash_trace def normal_function(): return "Hello, World!" @@ -53,6 +54,7 @@ def normal_method(self): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @codeflash_trace def normal_method(self): @@ -82,6 +84,7 @@ def class_method(cls): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @classmethod @codeflash_trace @@ -112,6 +115,7 @@ def static_method(): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @staticmethod @codeflash_trace @@ -141,6 +145,7 @@ def __init__(self, value): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @codeflash_trace def __init__(self, value): @@ -171,6 +176,7 @@ def property_method(self): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @property @other_decorator @@ -205,6 +211,7 @@ def test_method(self): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @codeflash_trace def test_method(self): From 896aa52c0d738b9c19e965819dd0361200ded1c5 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 19 Mar 2025 15:04:47 -0700 Subject: [PATCH 012/122] modified printing of results, handle errors when collecting benchmarks --- codeflash/benchmarking/codeflash_trace.py | 4 - codeflash/benchmarking/get_trace_info.py | 34 +++++--- .../pytest_new_process_trace_benchmarks.py | 11 +-- codeflash/benchmarking/replay_test.py | 13 +-- codeflash/benchmarking/trace_benchmarks.py | 25 +++++- codeflash/benchmarking/utils.py | 24 ++++-- codeflash/discovery/functions_to_optimize.py | 9 +- codeflash/models/models.py | 1 + codeflash/optimization/function_optimizer.py | 45 +++------- codeflash/optimization/optimizer.py | 83 +++++++++---------- codeflash/result/explanation.py | 41 +++++++-- codeflash/verification/test_runner.py | 2 + 12 files changed, 164 insertions(+), 128 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 9b9afead7..14505efee 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -84,9 +84,5 @@ def wrapper(*args, **kwargs): return result return wrapper - - - - # Create a singleton instance codeflash_trace = CodeflashTrace() diff --git a/codeflash/benchmarking/get_trace_info.py b/codeflash/benchmarking/get_trace_info.py index e9a050b84..d43327af7 100644 --- a/codeflash/benchmarking/get_trace_info.py +++ b/codeflash/benchmarking/get_trace_info.py @@ -4,13 +4,11 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize -def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, float]]: +def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, int]]: """Process the trace file and extract timing data for all functions. Args: trace_path: Path to the trace file - all_functions_to_optimize: List of FunctionToOptimize objects (not used directly, - but kept for backward compatibility) Returns: A nested dictionary where: @@ -30,8 +28,7 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, floa # Query the function_calls table for all function calls cursor.execute( "SELECT module_name, class_name, function_name, " - "benchmark_file_name, benchmark_function_name, benchmark_line_number, " - "(time_ns - overhead_time_ns) as actual_time_ns " + "benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns " "FROM function_calls" ) @@ -66,7 +63,7 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, floa return result -def get_benchmark_timings(trace_path: Path) -> dict[str, float]: +def get_benchmark_timings(trace_path: Path) -> dict[str, int]: """Extract total benchmark timings from trace files. Args: @@ -75,32 +72,47 @@ def get_benchmark_timings(trace_path: Path) -> dict[str, float]: Returns: A dictionary mapping where: - Keys are benchmark filename :: benchmark test function :: line number - - Values are total benchmark timing in milliseconds + - Values are total benchmark timing in milliseconds (with overhead subtracted) """ # Initialize the result dictionary result = {} + overhead_by_benchmark = {} # Connect to the SQLite database connection = sqlite3.connect(trace_path) cursor = connection.cursor() try: - # Query the benchmark_timings table + # Query the function_calls table to get total overhead for each benchmark + cursor.execute( + "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " + "FROM function_calls " + "GROUP BY benchmark_file_name, 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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + 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_name, benchmark_function_name, benchmark_line_number, time_ns " "FROM benchmark_timings" ) - # Process each row + # Process each row and subtract overhead for row in cursor.fetchall(): benchmark_file, benchmark_func, benchmark_line, time_ns = row # Create the benchmark key (file::function::line) benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - # Store the timing - result[benchmark_key] = time_ns + # Subtract overhead from total time + overhead = overhead_by_benchmark.get(benchmark_key, 0) + result[benchmark_key] = time_ns - overhead finally: # Close the connection diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index a83196758..6d4c85f41 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -20,13 +20,14 @@ db.setup() exitcode = pytest.main( [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashBenchmarkPlugin()] - ) + ) # Errors will be printed to stdout, not stderr db.write_function_timings(codeflash_trace.function_calls_data) db.write_benchmark_timings(CodeFlashBenchmarkPlugin.benchmark_timings) - db.print_function_timings() - db.print_benchmark_timings() + # db.print_function_timings() + # db.print_benchmark_timings() db.close() except Exception as e: - print(f"Failed to collect tests: {e!s}") - exitcode = -1 \ No newline at end of file + print(f"Failed to collect tests: {e!s}", file=sys.stderr) + exitcode = -1 + sys.exit(exitcode) \ No newline at end of file diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 75ef7e96d..a1d5b370a 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -142,8 +142,6 @@ def create_trace_replay_test_code( class_name = func.get("class_name") file_name = func.get("file_name") function_properties = func.get("function_properties") - print(f"Class: {class_name}, Function: {function_name}") - print(function_properties) if not class_name: alias = get_function_alias(module_name, function_name) test_body = test_function_body.format( @@ -197,7 +195,7 @@ def create_trace_replay_test_code( return imports + "\n" + metadata + "\n" + test_template -def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework: str = "pytest", max_run_count: int = 100) -> None: +def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework: str = "pytest", max_run_count: int = 100) -> int: """Generate multiple replay tests from the traced function calls, grouping by benchmark name. Args: @@ -211,6 +209,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework Dictionary mapping benchmark names to generated test code """ + count = 0 try: # Connect to the database conn = sqlite3.connect(trace_file_path.as_posix()) @@ -253,7 +252,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework }) if not functions_data: - print(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") + logger.info(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") continue # Generate the test code for this benchmark @@ -273,9 +272,11 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework # Write test code to file, parents = true output_dir.mkdir(parents=True, exist_ok=True) output_file.write_text(test_code, "utf-8") - print(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") + count += 1 + logger.info(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") conn.close() except Exception as e: - print(f"Error generating replay tests: {e}") + logger.info(f"Error generating replay tests: {e}") + return count \ No newline at end of file diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 9ae69495d..79395db79 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -1,10 +1,15 @@ from __future__ import annotations + +import re + +from pytest import ExitCode + +from codeflash.cli_cmds.console import logger from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE from pathlib import Path import subprocess def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, trace_file: Path) -> None: - # set up .trace databases result = subprocess.run( [ SAFE_SYS_EXECUTABLE, @@ -19,5 +24,19 @@ def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root text=True, env={"PYTHONPATH": str(project_root)}, ) - print("stdout:", result.stdout) - print("stderr:", result.stderr) + if result.returncode != 0: + if "ERROR collecting" in result.stdout: + # Pattern matches "===== ERRORS =====" (any number of =) and captures everything after + error_pattern = r"={3,}\s*ERRORS\s*={3,}\n([\s\S]*?)(?:={3,}|$)" + match = re.search(error_pattern, result.stdout) + error_section = match.group(1) if match else result.stdout + elif "FAILURES" in result.stdout: + # Pattern matches "===== FAILURES =====" (any number of =) and captures everything after + error_pattern = r"={3,}\s*FAILURES\s*={3,}\n([\s\S]*?)(?:={3,}|$)" + match = re.search(error_pattern, result.stdout) + error_section = match.group(1) if match else result.stdout + else: + error_section = result.stdout + logger.warning( + f"Error collecting benchmarks - Pytest Exit code: {result.returncode}={ExitCode(result.returncode).name}\n {error_section}" + ) \ No newline at end of file diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index d97c2e36e..685bfe739 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,7 +1,12 @@ -def print_benchmark_table(function_benchmark_timings, total_benchmark_timings): +def print_benchmark_table(function_benchmark_timings: dict[str,dict[str,int]], total_benchmark_timings: dict[str,int]): + # Define column widths + benchmark_col_width = 50 + time_col_width = 15 + # Print table header - print(f"{'Benchmark Test':<50} | {'Total Time (s)':<15} | {'Function Time (s)':<15} | {'Percentage (%)':<15}") - print("-" * 100) + header = f"{'Benchmark Test':{benchmark_col_width}} | {'Total Time (ms)':{time_col_width}} | {'Function Time (ms)':{time_col_width}} | {'Percentage (%)':{time_col_width}}" + print(header) + print("-" * len(header)) # Process each function's benchmark data for func_path, test_times in function_benchmark_timings.items(): @@ -14,13 +19,16 @@ def print_benchmark_table(function_benchmark_timings, total_benchmark_timings): total_time = total_benchmark_timings.get(test_name, 0) if total_time > 0: percentage = (func_time / total_time) * 100 - sorted_tests.append((test_name, total_time, func_time, percentage)) + # Convert nanoseconds to milliseconds + func_time_ms = func_time / 1_000_000 + total_time_ms = total_time / 1_000_000 + sorted_tests.append((test_name, total_time_ms, func_time_ms, percentage)) sorted_tests.sort(key=lambda x: x[3], reverse=True) # Print each test's data for test_name, total_time, func_time, percentage in sorted_tests: - print(f"{test_name:<50} | {total_time:<15.3f} | {func_time:<15.3f} | {percentage:<15.2f}") - -# Usage - + benchmark_file, benchmark_func, benchmark_line = test_name.split("::") + benchmark_name = f"{benchmark_file}::{benchmark_func}" + print(f"{benchmark_name:{benchmark_col_width}} | {total_time:{time_col_width}.3f} | {func_time:{time_col_width}.3f} | {percentage:{time_col_width}.2f}") + print() diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index cd1e53f9b..fb80541aa 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -106,7 +106,7 @@ def generic_visit(self, node: ast.AST) -> None: @dataclass(frozen=True, config={"arbitrary_types_allowed": True}) class FunctionToOptimize: - """Represents a function that is a candidate for optimization. + """Represent a function that is a candidate for optimization. Attributes ---------- @@ -121,6 +121,7 @@ class FunctionToOptimize: method extends this with the module name from the project root. """ + function_name: str file_path: Path parents: list[FunctionParent] # list[ClassDef | FunctionDef | AsyncFunctionDef] @@ -143,12 +144,6 @@ def qualified_name(self) -> str: def qualified_name_with_modules_from_root(self, project_root_path: Path) -> str: return f"{module_name_from_file_path(self.file_path, project_root_path)}.{self.qualified_name}" - # - # @property - # def qualified_name_with_file_name(self) -> str: - # class_name = self.parents[0].name if self.parents else None - # return f"{self.file_path}:{(class_name + ':' if class_name else '')}{self.function_name}" - def get_functions_to_optimize( optimize_all: str | None, diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 6419a3033..8ca3607d5 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -72,6 +72,7 @@ class BestOptimization(BaseModel): helper_functions: list[FunctionSource] runtime: int replay_runtime: int | None + replay_performance_gain: float | None winning_behavioral_test_results: TestResults winning_benchmarking_test_results: TestResults winning_replay_benchmarking_test_results : TestResults | None = None diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index b8814b434..bffdb276b 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -87,8 +87,8 @@ def __init__( function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_ast: ast.FunctionDef | None = None, aiservice_client: AiServiceClient | None = None, - function_benchmark_timings: dict[str, dict[str, float]] | None = None, - total_benchmark_timings: dict[str, float] | None = None, + function_benchmark_timings: dict[str, dict[str, int]] | None = None, + total_benchmark_timings: dict[str, int] | None = None, args: Namespace | None = None, ) -> None: self.project_root = test_cfg.project_root_path @@ -271,30 +271,10 @@ def optimize_function(self) -> Result[BestOptimization, str]: best_runtime_ns=best_optimization.runtime, function_name=function_to_optimize_qualified_name, file_path=self.function_to_optimize.file_path, + replay_performance_gain=best_optimization.replay_performance_gain if self.args.benchmark else None, + fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)] if self.args.benchmark else None, + total_benchmark_timings = self.total_benchmark_timings if self.args.benchmark else None, ) - speedup = explanation.speedup # - if self.args.benchmark: - original_replay_timing = original_code_baseline.benchmarking_test_results.total_replay_test_runtime() - fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)] - for benchmark_key, og_benchmark_timing in fto_benchmark_timings.items(): - # benchmark key is benchmark filename :: benchmark test function :: line number - try: - benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") - except ValueError: - print(f"Benchmark key {benchmark_key} is not in the expected format.") - continue - print(f"Calculating speedup for benchmark {benchmark_key}") - total_benchmark_timing = self.total_benchmark_timings[benchmark_key] - # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values - print(f"Original benchmark timing: {total_benchmark_timing}") - replay_speedup = original_replay_timing / best_optimization.replay_runtime - 1 - print(f"Replay speedup: {replay_speedup}") - expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / (replay_speedup + 1) * og_benchmark_timing - print(f"Expected new benchmark timing: {expected_new_benchmark_timing}") - - benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing - benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 - print(f"Benchmark speedup: {benchmark_speedup_percent:.2f}%") self.log_successful_optimization(explanation, generated_tests) @@ -450,21 +430,21 @@ def determine_best_candidate( original_runtime_ns=original_code_replay_runtime, optimized_runtime_ns=candidate_replay_runtime, ) - tree.add("Replay Benchmarking： ") - tree.add(f"Original summed runtime: {humanize_runtime(original_code_replay_runtime)}") + tree.add(f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") tree.add( - f"Best summed runtime: {humanize_runtime(candidate_replay_runtime)} " + f"Best benchmark replay runtime: {humanize_runtime(candidate_replay_runtime)} " f"(measured over {candidate_result.max_loop_count} " f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" ) - tree.add(f"Speedup percentage: {replay_perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio: {replay_perf_gain + 1:.1f}X") + tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain + 1:.1f}X") best_optimization = BestOptimization( candidate=candidate, helper_functions=code_context.helper_functions, runtime=best_test_runtime, replay_runtime=candidate_replay_runtime if self.args.benchmark else None, winning_behavioral_test_results=candidate_result.behavior_test_results, + replay_performance_gain=replay_perf_gain if self.args.benchmark else None, winning_benchmarking_test_results=candidate_result.benchmarking_test_results, winning_replay_benchmarking_test_results=candidate_result.benchmarking_test_results, ) @@ -520,7 +500,8 @@ def log_successful_optimization(self, explanation: Explanation, generated_tests: ) console.print(Group(explanation_panel, tests_panel)) - console.print(explanation_panel) + else: + console.print(explanation_panel) ph( "cli-optimize-success", @@ -633,7 +614,6 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi existing_test_files_count += 1 elif test_type == TestType.REPLAY_TEST: replay_test_files_count += 1 - print("Replay test found") elif test_type == TestType.CONCOLIC_COVERAGE_TEST: concolic_coverage_test_files_count += 1 else: @@ -1107,7 +1087,6 @@ def run_and_parse_tests( f"stdout: {run_result.stdout}\n" f"stderr: {run_result.stderr}\n" ) - # print(test_files) results, coverage_results = parse_test_results( test_xml_path=result_file_path, test_files=test_files, diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 953f0b755..848960cf3 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -11,8 +11,9 @@ from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.utils import print_benchmark_table -from codeflash.cli_cmds.console import console, logger +from codeflash.cli_cmds.console import console, logger, progress_bar from codeflash.code_utils import env_utils +from codeflash.code_utils.code_extractor import add_needed_imports_from_module from codeflash.code_utils.code_replacer import normalize_code, normalize_node from codeflash.code_utils.code_utils import get_run_tmp_file from codeflash.code_utils.static_analysis import analyze_imported_modules, get_first_top_level_function_or_method_ast @@ -96,42 +97,43 @@ def run(self) -> None: project_root=self.args.project_root, module_root=self.args.module_root, ) - all_functions_to_optimize = [ - fto for functions_to_optimize in file_to_funcs_to_optimize.values() for fto in functions_to_optimize] + function_benchmark_timings = None + total_benchmark_timings = None if self.args.benchmark: - # 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: - instrument_codeflash_trace_decorator(fto) - trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" - trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Simply run all tests that use pytest-benchmark - generate_replay_test(trace_file, Path(self.args.tests_root) / "codeflash_replay_tests" ) - function_benchmark_timings = get_function_benchmark_timings(trace_file) - total_benchmark_timings = get_benchmark_timings(trace_file) - print(function_benchmark_timings) - print(total_benchmark_timings) - logger.info("Finished tracing existing benchmarks") - except Exception as e: - logger.info(f"Error while tracing existing benchmarks: {e}") - logger.info(f"Information on existing benchmarks will not be available for this run.") - 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]) - + with progress_bar( + f"Running benchmarks in {self.args.benchmarks_root}", + transient=True, + ): + # 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: + instrument_codeflash_trace_decorator(fto) + trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" + replay_tests_dir = Path(self.args.tests_root) / "codeflash_replay_tests" + trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Simply run all tests that use pytest-benchmark + replay_count = generate_replay_test(trace_file, replay_tests_dir) + if replay_count == 0: + logger.info(f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization") + else: + function_benchmark_timings = get_function_benchmark_timings(trace_file) + total_benchmark_timings = get_benchmark_timings(trace_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) + print_benchmark_table(function_benchmark_timings, total_benchmark_timings) + logger.info("Finished tracing existing benchmarks") + except Exception as e: + logger.info(f"Error while tracing existing benchmarks: {e}") + logger.info(f"Information on existing benchmarks will not be available for this run.") + 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]) - # return optimizations_found: int = 0 function_iterator_count: int = 0 if self.args.test_framework == "pytest": @@ -211,15 +213,10 @@ def run(self) -> None: f"Skipping optimization." ) continue - if self.args.benchmark: - + if self.args.benchmark and function_benchmark_timings and total_benchmark_timings: function_optimizer = self.create_function_optimizer( function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code, function_benchmark_timings, total_benchmark_timings ) - # function_optimizer = self.create_function_optimizer( - # function_to_optimize, function_to_optimize_ast, function_to_tests, - # validated_original_code[original_module_path].source_code - # ) else: function_optimizer = self.create_function_optimizer( function_to_optimize, function_to_optimize_ast, function_to_tests, @@ -251,9 +248,9 @@ def run(self) -> None: if function_optimizer.test_cfg.concolic_test_root_dir: shutil.rmtree(function_optimizer.test_cfg.concolic_test_root_dir, ignore_errors=True) if self.args.benchmark: - trace_dir = Path(self.args.benchmarks_root) / "codeflash_replay_tests" - if trace_dir.exists(): - shutil.rmtree(trace_dir, ignore_errors=True) + if replay_tests_dir.exists(): + shutil.rmtree(replay_tests_dir, ignore_errors=True) + trace_file.unlink(missing_ok=True) if hasattr(get_run_tmp_file, "tmpdir"): get_run_tmp_file.tmpdir.cleanup() diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 1dd53ceb5..4f3badb58 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -15,6 +15,9 @@ class Explanation: best_runtime_ns: int function_name: str file_path: Path + replay_performance_gain: float | None + fto_benchmark_timings: dict[str, int] | None + total_benchmark_timings: dict[str, int] | None @property def perf_improvement_line(self) -> str: @@ -37,16 +40,38 @@ def to_console_string(self) -> str: # TODO: Sometimes the explanation says something similar to "This is the code that was optimized", remove such parts original_runtime_human = humanize_runtime(self.original_runtime_ns) best_runtime_human = humanize_runtime(self.best_runtime_ns) + benchmark_info = "" + if self.replay_performance_gain: + benchmark_info += "Benchmark Performance Details:\n" + for benchmark_key, og_benchmark_timing in self.fto_benchmark_timings.items(): + # benchmark key is benchmark filename :: benchmark test function :: line number + try: + benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") + except ValueError: + benchmark_info += f"Benchmark key {benchmark_key} is not in the expected format.\n" + continue + + total_benchmark_timing = self.total_benchmark_timings[benchmark_key] + # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values + benchmark_info += f"Original timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(total_benchmark_timing)}\n" + replay_speedup = self.replay_performance_gain + expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / ( + replay_speedup + 1) * og_benchmark_timing + benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(expected_new_benchmark_timing)}\n" + + benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing + benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 + benchmark_info += f"Benchmark speedup for {benchmark_file_name}::{benchmark_test_function}: {benchmark_speedup_percent:.2f}%\n\n" return ( - f"Optimized {self.function_name} in {self.file_path}\n" - f"{self.perf_improvement_line}\n" - f"Runtime went down from {original_runtime_human} to {best_runtime_human} \n\n" - + "Explanation:\n" - + self.raw_explanation_message - + " \n\n" - + "The new optimized code was tested for correctness. The results are listed below.\n" - + f"{TestResults.report_to_string(self.winning_behavioral_test_results.get_test_pass_fail_report_by_type())}\n" + f"Optimized {self.function_name} in {self.file_path}\n" + f"{self.perf_improvement_line}\n" + f"Runtime went down from {original_runtime_human} to {best_runtime_human} \n\n" + + (benchmark_info if benchmark_info else "") + + self.raw_explanation_message + + " \n\n" + + "The new optimized code was tested for correctness. The results are listed below.\n" + + f"{TestResults.report_to_string(self.winning_behavioral_test_results.get_test_pass_fail_report_by_type())}\n" ) def explanation_message(self) -> str: diff --git a/codeflash/verification/test_runner.py b/codeflash/verification/test_runner.py index 578add352..6e6c0990d 100644 --- a/codeflash/verification/test_runner.py +++ b/codeflash/verification/test_runner.py @@ -57,6 +57,7 @@ def run_behavioral_tests( ) test_files = list(set(test_files)) # remove multiple calls in the same test function common_pytest_args = [ + "--benchmark-skip", "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", @@ -160,6 +161,7 @@ def run_benchmarking_tests( test_files.append(str(file.benchmarking_file_path)) test_files = list(set(test_files)) # remove multiple calls in the same test function pytest_args = [ + "--benchmark-skip", "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", From ad17de464061e5f13e1a3bb98721faa407f7fab2 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 19 Mar 2025 15:59:13 -0700 Subject: [PATCH 013/122] tests pass --- code_to_optimize/bubble_sort.py | 2 ++ ...process_and_bubble_sort_codeflash_trace.py | 28 +++++++++++++++ codeflash/discovery/functions_to_optimize.py | 34 ++++++++++--------- .../discovery/pytest_new_process_discovery.py | 2 +- tests/test_unit_test_discovery.py | 9 ++--- 5 files changed, 52 insertions(+), 23 deletions(-) create mode 100644 code_to_optimize/process_and_bubble_sort_codeflash_trace.py diff --git a/code_to_optimize/bubble_sort.py b/code_to_optimize/bubble_sort.py index db7db5f92..9e97f63a0 100644 --- a/code_to_optimize/bubble_sort.py +++ b/code_to_optimize/bubble_sort.py @@ -1,8 +1,10 @@ def sorter(arr): + print("codeflash stdout: Sorting list") for i in range(len(arr)): for j in range(len(arr) - 1): if arr[j] > arr[j + 1]: temp = arr[j] arr[j] = arr[j + 1] arr[j + 1] = temp + print(f"result: {arr}") return arr diff --git a/code_to_optimize/process_and_bubble_sort_codeflash_trace.py b/code_to_optimize/process_and_bubble_sort_codeflash_trace.py new file mode 100644 index 000000000..37c2abab8 --- /dev/null +++ b/code_to_optimize/process_and_bubble_sort_codeflash_trace.py @@ -0,0 +1,28 @@ +from code_to_optimize.bubble_sort import sorter +from codeflash.benchmarking.codeflash_trace import codeflash_trace + +def calculate_pairwise_products(arr): + """ + Calculate the average of all pairwise products in the array. + """ + sum_of_products = 0 + count = 0 + + for i in range(len(arr)): + for j in range(len(arr)): + if i != j: + sum_of_products += arr[i] * arr[j] + count += 1 + + # The average of all pairwise products + return sum_of_products / count if count > 0 else 0 + +@codeflash_trace +def compute_and_sort(arr): + # Compute pairwise sums average + pairwise_average = calculate_pairwise_products(arr) + + # Call sorter function + sorter(arr.copy()) + + return pairwise_average diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index fb80541aa..cd0bfc50a 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -363,23 +363,25 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None: for decorator in body_node.decorator_list ): self.is_staticmethod = True + print(f"static method found: {self.function_name}") + return + elif self.line_no: + # If we have line number info, check if class has a static method with the same line number + # This way, if we don't have the class name, we can still find the static method + for body_node in node.body: + if ( + isinstance(body_node, ast.FunctionDef) + and body_node.name == self.function_name + and body_node.lineno in {self.line_no, self.line_no + 1} + and any( + isinstance(decorator, ast.Name) and decorator.id == "staticmethod" + for decorator in body_node.decorator_list + ) + ): + self.is_staticmethod = True + self.is_top_level = True + self.class_name = node.name return - # else: - # # search if the class has a staticmethod with the same name and on the same line number - # for body_node in node.body: - # if ( - # isinstance(body_node, ast.FunctionDef) - # and body_node.name == self.function_name - # # and body_node.lineno in {self.line_no, self.line_no + 1} - # and any( - # isinstance(decorator, ast.Name) and decorator.id == "staticmethod" - # for decorator in body_node.decorator_list - # ) - # ): - # self.is_staticmethod = True - # self.is_top_level = True - # self.class_name = node.name - # return return diff --git a/codeflash/discovery/pytest_new_process_discovery.py b/codeflash/discovery/pytest_new_process_discovery.py index 397eabe01..d5a80f501 100644 --- a/codeflash/discovery/pytest_new_process_discovery.py +++ b/codeflash/discovery/pytest_new_process_discovery.py @@ -34,7 +34,7 @@ def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, s try: exitcode = pytest.main( - [tests_root, "-pno:logging", "--collect-only", "-m", "not skip"], plugins=[PytestCollectionPlugin()] + [tests_root, "-pno:logging", "--collect-only", "-m", "not skip", "--benchmark-skip"], plugins=[PytestCollectionPlugin()] ) except Exception as e: # noqa: BLE001 print(f"Failed to collect tests: {e!s}") # noqa: T201 diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index fe56b907f..ba658f46e 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -18,11 +18,11 @@ def test_unit_test_discovery_pytest(): ) tests = discover_unit_tests(test_config) assert len(tests) > 0 - # print(tests) + def test_benchmark_test_discovery_pytest(): project_path = Path(__file__).parent.parent.resolve() / "code_to_optimize" - tests_path = project_path / "tests" / "pytest" / "benchmarks" / "test_benchmark_bubble_sort.py" + tests_path = project_path / "tests" / "pytest" / "benchmarks" test_config = TestConfig( tests_root=tests_path, project_root_path=project_path, @@ -30,10 +30,7 @@ def test_benchmark_test_discovery_pytest(): tests_project_rootdir=tests_path.parent, ) tests = discover_unit_tests(test_config) - assert len(tests) > 0 - assert 'bubble_sort.sorter' in tests - benchmark_tests = sum(1 for test in tests['bubble_sort.sorter'] if test.tests_in_file.test_type == TestType.BENCHMARK_TEST) - assert benchmark_tests == 1 + assert len(tests) == 1 # Should not discover benchmark tests def test_unit_test_discovery_unittest(): From 92e6bf5981e591d60d5c501ba4d153719a68c6da Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 19 Mar 2025 16:00:45 -0700 Subject: [PATCH 014/122] revert pyproject.toml --- pyproject.toml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 877815004..2e71f2a0a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -216,9 +216,8 @@ initial-content = """ [tool.codeflash] -module-root = "code_to_optimize" -tests-root = "code_to_optimize/tests" -benchmarks-root = "code_to_optimize/tests/pytest/benchmarks" +module-root = "codeflash" +tests-root = "tests" test-framework = "pytest" formatter-cmds = [ "uvx ruff check --exit-zero --fix $file", From 4784723b5664fa338f978b8f22ebff4402bf555b Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 19 Mar 2025 16:02:44 -0700 Subject: [PATCH 015/122] mypy fixes --- codeflash/result/explanation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 4f3badb58..dd74ddbe6 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -41,7 +41,7 @@ def to_console_string(self) -> str: original_runtime_human = humanize_runtime(self.original_runtime_ns) best_runtime_human = humanize_runtime(self.best_runtime_ns) benchmark_info = "" - if self.replay_performance_gain: + if self.replay_performance_gain and self.fto_benchmark_timings and self.total_benchmark_timings: benchmark_info += "Benchmark Performance Details:\n" for benchmark_key, og_benchmark_timing in self.fto_benchmark_timings.items(): # benchmark key is benchmark filename :: benchmark test function :: line number @@ -57,7 +57,7 @@ def to_console_string(self) -> str: replay_speedup = self.replay_performance_gain expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / ( replay_speedup + 1) * og_benchmark_timing - benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(expected_new_benchmark_timing)}\n" + benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(int(expected_new_benchmark_timing))}\n" benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 From b77a979de0ad78024ec33708bfaa65c70bc4291c Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 10:10:59 -0700 Subject: [PATCH 016/122] import changes --- codeflash/benchmarking/codeflash_trace.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 14505efee..b1236ffbf 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -1,9 +1,7 @@ import functools import os import pickle -import sqlite3 import time -from pathlib import Path from typing import Callable From 0c2a3b6c693ab213b1c6251333dbf7d8839e4a5f Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 13:45:36 -0700 Subject: [PATCH 017/122] removed benchmark skip command --- codeflash/verification/test_runner.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/codeflash/verification/test_runner.py b/codeflash/verification/test_runner.py index 190b14620..b5a31f0d7 100644 --- a/codeflash/verification/test_runner.py +++ b/codeflash/verification/test_runner.py @@ -59,7 +59,6 @@ def run_behavioral_tests( ) test_files = list(set(test_files)) # remove multiple calls in the same test function common_pytest_args = [ - "--benchmark-skip", "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", @@ -165,7 +164,6 @@ def run_benchmarking_tests( test_files.append(str(file.benchmarking_file_path)) test_files = list(set(test_files)) # remove multiple calls in the same test function pytest_args = [ - "--benchmark-skip", "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", From 9a41bdd93132ab09fdea0c1ada92cab294105ba1 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 15:09:21 -0700 Subject: [PATCH 018/122] shifted benchmark class in plugin, improved display of benchmark info --- codeflash/benchmarking/plugin/plugin.py | 44 ++++++++++++++----------- codeflash/benchmarking/utils.py | 37 ++++++++++++++------- 2 files changed, 49 insertions(+), 32 deletions(-) diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index a5f82fc3a..ee7504ec4 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -5,6 +5,29 @@ import os class CodeFlashBenchmarkPlugin: benchmark_timings = [] + + class Benchmark: + def __init__(self, request): + self.request = request + + def __call__(self, func, *args, **kwargs): + benchmark_file_name = self.request.node.fspath.basename + benchmark_function_name = self.request.node.name + line_number = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack + + os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name + os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = benchmark_file_name + os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = line_number + os.environ["CODEFLASH_BENCHMARKING"] = "True" + + start = time.perf_counter_ns() + result = func(*args, **kwargs) + end = time.perf_counter_ns() + + os.environ["CODEFLASH_BENCHMARKING"] = "False" + CodeFlashBenchmarkPlugin.benchmark_timings.append( + (benchmark_file_name, benchmark_function_name, line_number, end - start)) + return result @staticmethod def pytest_addoption(parser): parser.addoption( @@ -36,23 +59,4 @@ def benchmark(request): if not request.config.getoption("--codeflash-trace"): return None - class Benchmark: - - def __call__(self, func, *args, **kwargs): - benchmark_file_name = request.node.fspath.basename - benchmark_function_name = request.node.name - line_number = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack - os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name - os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = benchmark_file_name - os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = line_number - os.environ["CODEFLASH_BENCHMARKING"] = "True" - - start = time.perf_counter_ns() - result = func(*args, **kwargs) - end = time.perf_counter_ns() - - os.environ["CODEFLASH_BENCHMARKING"] = "False" - CodeFlashBenchmarkPlugin.benchmark_timings.append((benchmark_file_name, benchmark_function_name, line_number, end - start)) - return result - - return Benchmark() + return CodeFlashBenchmarkPlugin.Benchmark(request) \ No newline at end of file diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 685bfe739..becf606a4 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,17 +1,23 @@ -def print_benchmark_table(function_benchmark_timings: dict[str,dict[str,int]], total_benchmark_timings: dict[str,int]): - # Define column widths - benchmark_col_width = 50 - time_col_width = 15 +from rich.console import Console +from rich.table import Table - # Print table header - header = f"{'Benchmark Test':{benchmark_col_width}} | {'Total Time (ms)':{time_col_width}} | {'Function Time (ms)':{time_col_width}} | {'Percentage (%)':{time_col_width}}" - print(header) - print("-" * len(header)) + +def print_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]], + total_benchmark_timings: dict[str, int]): + console = Console() # Process each function's benchmark data for func_path, test_times in function_benchmark_timings.items(): function_name = func_path.split(":")[-1] - print(f"\n== Function: {function_name} ==") + + # 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) + 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") # Sort by percentage (highest first) sorted_tests = [] @@ -26,9 +32,16 @@ def print_benchmark_table(function_benchmark_timings: dict[str,dict[str,int]], t sorted_tests.sort(key=lambda x: x[3], reverse=True) - # Print each test's data + # Add rows to the table for test_name, total_time, func_time, percentage in sorted_tests: benchmark_file, benchmark_func, benchmark_line = test_name.split("::") benchmark_name = f"{benchmark_file}::{benchmark_func}" - print(f"{benchmark_name:{benchmark_col_width}} | {total_time:{time_col_width}.3f} | {func_time:{time_col_width}.3f} | {percentage:{time_col_width}.2f}") - print() + table.add_row( + benchmark_name, + f"{total_time:.3f}", + f"{func_time:.3f}", + f"{percentage:.2f}" + ) + + # Print the table + console.print(table) \ No newline at end of file From 5577cd539b979098d05444f3f92e070d48221a43 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 15:11:17 -0700 Subject: [PATCH 019/122] cleanup tests better --- tests/test_trace_benchmarks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index c49e7c693..4bc0ef278 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -150,5 +150,5 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): assert test_sort_path.read_text("utf-8").strip()==test_sort_code.strip() finally: # cleanup - shutil.rmtree(tests_root) - pass \ No newline at end of file + if tests_root.exists(): + shutil.rmtree(tests_root, ignore_errors=True) \ No newline at end of file From 80730f956963c655c128792a060b21291e93200a Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 15:49:26 -0700 Subject: [PATCH 020/122] modified paths in test --- tests/test_trace_benchmarks.py | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 4bc0ef278..4113d954e 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -32,34 +32,36 @@ def test_trace_benchmarks(): # Assert the length of function calls assert len(function_calls) == 7, f"Expected 6 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() # Expected function calls expected_calls = [ ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_class_sort", "test_benchmark_bubble_sort.py", 20), ("sort_class", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_class_sort", "test_benchmark_bubble_sort.py", 18), ("sort_static", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_class_sort", "test_benchmark_bubble_sort.py", 19), ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_class_sort", "test_benchmark_bubble_sort.py", 17), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_sort", "test_benchmark_bubble_sort.py", 7), ("compute_and_sort", "", "code_to_optimize.process_and_bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/process_and_bubble_sort_codeflash_trace.py'}", + f"{process_and_bubble_sort_path}", "test_compute_and_sort", "test_process_and_sort.py", 4), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_no_func", "test_process_and_sort.py", 8), ] for idx, (actual, expected) in enumerate(zip(function_calls, expected_calls)): @@ -86,7 +88,7 @@ def test_trace_benchmarks(): trace_file_path = r"{output_file.as_posix()}" def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "sorter" @@ -99,7 +101,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): ret = instance.sorter(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) if not args: @@ -107,13 +109,13 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_class(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_static(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_static(*args, **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "__init__" @@ -141,7 +143,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): trace_file_path = r"{output_file}" def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"{bubble_sort_path}", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_sorter(*args, **kwargs) @@ -150,5 +152,5 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): assert test_sort_path.read_text("utf-8").strip()==test_sort_code.strip() finally: # cleanup - if tests_root.exists(): - shutil.rmtree(tests_root, ignore_errors=True) \ No newline at end of file + shutil.rmtree(tests_root) + pass \ No newline at end of file From d610f8cc114cb10d4b9697ab5511eafc7fbf978e Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 15:53:17 -0700 Subject: [PATCH 021/122] typing fix --- codeflash/models/models.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 8ca3607d5..40e1f3906 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -71,11 +71,11 @@ class BestOptimization(BaseModel): candidate: OptimizedCandidate helper_functions: list[FunctionSource] runtime: int - replay_runtime: int | None - replay_performance_gain: float | None + replay_runtime: Optional[int] = None + replay_performance_gain: Optional[float] = None winning_behavioral_test_results: TestResults winning_benchmarking_test_results: TestResults - winning_replay_benchmarking_test_results : TestResults | None = None + winning_replay_benchmarking_test_results : Optional[TestResults] = None class CodeString(BaseModel): @@ -220,7 +220,7 @@ class OriginalCodeBaseline(BaseModel): behavioral_test_results: TestResults benchmarking_test_results: TestResults runtime: int - coverage_results: CoverageData | None + coverage_results: Optional[CoverageData] class CoverageStatus(Enum): From 93f583c7c2f9ddf0d3ae073ecba382a48a9a5e96 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 17:01:29 -0700 Subject: [PATCH 022/122] typing fix for 3.9 --- codeflash/result/explanation.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index dd74ddbe6..434683417 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -1,4 +1,6 @@ +from __future__ import annotations from pathlib import Path +from typing import Optional, Union from pydantic.dataclasses import dataclass @@ -15,9 +17,9 @@ class Explanation: best_runtime_ns: int function_name: str file_path: Path - replay_performance_gain: float | None - fto_benchmark_timings: dict[str, int] | None - total_benchmark_timings: dict[str, int] | None + replay_performance_gain: Optional[float] + fto_benchmark_timings: Optional[Union[dict, int]] + total_benchmark_timings: Optional[Union[dict, int]] @property def perf_improvement_line(self) -> str: From d422e354f2ed4953bd46c8b2eead806f0b0a1345 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 21 Mar 2025 13:09:28 -0700 Subject: [PATCH 023/122] typing fix for 3.9 --- codeflash/result/explanation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 434683417..248417452 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -18,8 +18,8 @@ class Explanation: function_name: str file_path: Path replay_performance_gain: Optional[float] - fto_benchmark_timings: Optional[Union[dict, int]] - total_benchmark_timings: Optional[Union[dict, int]] + fto_benchmark_timings: Optional[dict[str, int]] + total_benchmark_timings: Optional[dict[str, int]] @property def perf_improvement_line(self) -> str: From 163781081a874b9e08e6480403cb72ef4da10566 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Mon, 24 Mar 2025 16:45:13 -0700 Subject: [PATCH 024/122] works with multithreading, added test --- code_to_optimize/bubble_sort_multithread.py | 23 +++++++ .../benchmarks/test_benchmark_bubble_sort.py | 9 +-- .../benchmarks/test_process_and_sort.py | 4 +- .../test_multithread_sort.py | 4 ++ .../test_benchmark_bubble_sort.py | 20 ++++++ .../benchmarks_test/test_process_and_sort.py | 8 +++ codeflash/benchmarking/codeflash_trace.py | 16 ++--- codeflash/benchmarking/utils.py | 61 ++++++++++++------- codeflash/optimization/function_optimizer.py | 4 +- codeflash/optimization/optimizer.py | 16 ++--- codeflash/result/explanation.py | 22 ++++--- tests/test_trace_benchmarks.py | 61 ++++++++++++++++++- 12 files changed, 185 insertions(+), 63 deletions(-) create mode 100644 code_to_optimize/bubble_sort_multithread.py create mode 100644 code_to_optimize/tests/pytest/benchmarks_multithread/test_multithread_sort.py create mode 100644 code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py create mode 100644 code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py diff --git a/code_to_optimize/bubble_sort_multithread.py b/code_to_optimize/bubble_sort_multithread.py new file mode 100644 index 000000000..3659b01bf --- /dev/null +++ b/code_to_optimize/bubble_sort_multithread.py @@ -0,0 +1,23 @@ +# from code_to_optimize.bubble_sort_codeflash_trace import sorter +from code_to_optimize.bubble_sort_codeflash_trace import sorter +import concurrent.futures + + +def multithreaded_sorter(unsorted_lists: list[list[int]]) -> list[list[int]]: + # Create a list to store results in the correct order + sorted_lists = [None] * len(unsorted_lists) + + # Use ThreadPoolExecutor to manage threads + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + # Submit all sorting tasks and map them to their original indices + future_to_index = { + executor.submit(sorter, unsorted_list): i + for i, unsorted_list in enumerate(unsorted_lists) + } + + # Collect results as they complete + for future in concurrent.futures.as_completed(future_to_index): + index = future_to_index[future] + sorted_lists[index] = future.result() + + return sorted_lists \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py index 03b9d38d1..3d7b24a6c 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py @@ -1,6 +1,6 @@ import pytest -from code_to_optimize.bubble_sort_codeflash_trace import sorter, Sorter +from code_to_optimize.bubble_sort import sorter def test_sort(benchmark): @@ -11,10 +11,3 @@ def test_sort(benchmark): def test_sort2(): result = sorter(list(reversed(range(500)))) assert result == list(range(500)) - -def test_class_sort(benchmark): - obj = Sorter(list(reversed(range(100)))) - result1 = benchmark(obj.sorter, 2) - result2 = benchmark(Sorter.sort_class, list(reversed(range(100)))) - result3 = benchmark(Sorter.sort_static, list(reversed(range(100)))) - result4 = benchmark(Sorter, [1,2,3]) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py index bcd42eab9..8d31c926a 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py @@ -1,5 +1,5 @@ -from code_to_optimize.process_and_bubble_sort_codeflash_trace import compute_and_sort -from code_to_optimize.bubble_sort_codeflash_trace import sorter +from code_to_optimize.process_and_bubble_sort import compute_and_sort +from code_to_optimize.bubble_sort import sorter def test_compute_and_sort(benchmark): result = benchmark(compute_and_sort, list(reversed(range(500)))) assert result == 62208.5 diff --git a/code_to_optimize/tests/pytest/benchmarks_multithread/test_multithread_sort.py b/code_to_optimize/tests/pytest/benchmarks_multithread/test_multithread_sort.py new file mode 100644 index 000000000..4a5c68a2b --- /dev/null +++ b/code_to_optimize/tests/pytest/benchmarks_multithread/test_multithread_sort.py @@ -0,0 +1,4 @@ +from code_to_optimize.bubble_sort_multithread import multithreaded_sorter + +def test_benchmark_sort(benchmark): + benchmark(multithreaded_sorter, [list(range(1000)) for i in range (10)]) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py new file mode 100644 index 000000000..03b9d38d1 --- /dev/null +++ b/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py @@ -0,0 +1,20 @@ +import pytest + +from code_to_optimize.bubble_sort_codeflash_trace import sorter, Sorter + + +def test_sort(benchmark): + result = benchmark(sorter, list(reversed(range(500)))) + assert result == list(range(500)) + +# This should not be picked up as a benchmark test +def test_sort2(): + result = sorter(list(reversed(range(500)))) + assert result == list(range(500)) + +def test_class_sort(benchmark): + obj = Sorter(list(reversed(range(100)))) + result1 = benchmark(obj.sorter, 2) + result2 = benchmark(Sorter.sort_class, list(reversed(range(100)))) + result3 = benchmark(Sorter.sort_static, list(reversed(range(100)))) + result4 = benchmark(Sorter, [1,2,3]) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py new file mode 100644 index 000000000..bcd42eab9 --- /dev/null +++ b/code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py @@ -0,0 +1,8 @@ +from code_to_optimize.process_and_bubble_sort_codeflash_trace import compute_and_sort +from code_to_optimize.bubble_sort_codeflash_trace import sorter +def test_compute_and_sort(benchmark): + result = benchmark(compute_and_sort, list(reversed(range(500)))) + assert result == 62208.5 + +def test_no_func(benchmark): + benchmark(sorter, list(reversed(range(500)))) \ No newline at end of file diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index b1236ffbf..f708d752f 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -15,11 +15,6 @@ class CodeflashTrace: def __init__(self) -> None: self.function_calls_data = [] - # def __enter__(self) -> None: - # # Initialize for context manager use - # self.function_calls_data = [] - # return self - def __exit__(self, exc_type, exc_val, exc_tb) -> None: # Cleanup is optional here pass @@ -37,15 +32,14 @@ def __call__(self, func: Callable) -> Callable: @functools.wraps(func) def wrapper(*args, **kwargs): # Measure execution time - start_time = time.perf_counter_ns() + start_time = time.thread_time_ns() result = func(*args, **kwargs) - end_time = time.perf_counter_ns() - + end_time = time.thread_time_ns() # Calculate execution time execution_time = end_time - start_time # Measure overhead - overhead_start_time = time.perf_counter_ns() + overhead_start_time = time.thread_time_ns() try: # Check if currently in pytest benchmark fixture @@ -66,7 +60,7 @@ def wrapper(*args, **kwargs): if "." in qualname: class_name = qualname.split(".")[0] # Calculate overhead time - overhead_end_time = time.perf_counter_ns() + overhead_end_time = time.thread_time_ns() overhead_time = overhead_end_time - overhead_start_time @@ -75,7 +69,7 @@ def wrapper(*args, **kwargs): benchmark_function_name, benchmark_file_name, benchmark_line_number, execution_time, overhead_time, pickled_args, pickled_kwargs) ) - + print("appended") except Exception as e: print(f"Error in codeflash_trace: {e}") diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index becf606a4..eeacb6975 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,47 +1,64 @@ from rich.console import Console from rich.table import Table +from codeflash.cli_cmds.console import logger -def print_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]], - total_benchmark_timings: dict[str, int]): - console = Console() +def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]], + total_benchmark_timings: dict[str, int]) -> dict[str, list[tuple[str, float, float, float]]]: + function_to_result = {} # Process each function's benchmark data for func_path, test_times in function_benchmark_timings.items(): - 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) - 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") - # Sort by percentage (highest first) sorted_tests = [] for test_name, func_time in test_times.items(): total_time = total_benchmark_timings.get(test_name, 0) + if func_time > total_time: + logger.debug(f"Skipping test {test_name} due to func_time {func_time} > total_time {total_time}") + # If the function time is greater than total time, likely to have multithreading / multiprocessing issues. + # Do not try to project the optimization impact for this function. + sorted_tests.append((test_name, 0.0, 0.0, 0.0)) if total_time > 0: percentage = (func_time / total_time) * 100 # Convert nanoseconds to milliseconds func_time_ms = func_time / 1_000_000 total_time_ms = total_time / 1_000_000 sorted_tests.append((test_name, total_time_ms, func_time_ms, percentage)) - sorted_tests.sort(key=lambda x: x[3], reverse=True) + function_to_result[func_path] = sorted_tests + return function_to_result + +def print_benchmark_table(function_to_results: dict[str, list[tuple[str, float, float, float]]]) -> None: + console = Console() + for func_path, sorted_tests in function_to_results.items(): + 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) + 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") - # Add rows to the table for test_name, total_time, func_time, percentage in sorted_tests: benchmark_file, benchmark_func, benchmark_line = test_name.split("::") benchmark_name = f"{benchmark_file}::{benchmark_func}" - table.add_row( - benchmark_name, - f"{total_time:.3f}", - f"{func_time:.3f}", - f"{percentage:.2f}" - ) + if total_time == 0.0: + table.add_row( + benchmark_name, + "N/A", + "N/A", + "N/A" + ) + else: + table.add_row( + benchmark_name, + f"{total_time:.3f}", + f"{func_time:.3f}", + f"{percentage:.2f}" + ) # Print the table console.print(table) \ No newline at end of file diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index bffdb276b..70322c9fe 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -87,7 +87,7 @@ def __init__( function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_ast: ast.FunctionDef | None = None, aiservice_client: AiServiceClient | None = None, - function_benchmark_timings: dict[str, dict[str, int]] | None = None, + function_benchmark_timings: dict[str, int] | None = None, total_benchmark_timings: dict[str, int] | None = None, args: Namespace | None = None, ) -> None: @@ -272,7 +272,7 @@ def optimize_function(self) -> Result[BestOptimization, str]: function_name=function_to_optimize_qualified_name, file_path=self.function_to_optimize.file_path, replay_performance_gain=best_optimization.replay_performance_gain if self.args.benchmark else None, - fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)] if self.args.benchmark else None, + fto_benchmark_timings = self.function_benchmark_timings if self.args.benchmark else None, total_benchmark_timings = self.total_benchmark_timings if self.args.benchmark else None, ) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 848960cf3..4a43a4c0b 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -10,10 +10,9 @@ from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest -from codeflash.benchmarking.utils import print_benchmark_table +from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table from codeflash.cli_cmds.console import console, logger, progress_bar from codeflash.code_utils import env_utils -from codeflash.code_utils.code_extractor import add_needed_imports_from_module from codeflash.code_utils.code_replacer import normalize_code, normalize_node from codeflash.code_utils.code_utils import get_run_tmp_file from codeflash.code_utils.static_analysis import analyze_imported_modules, get_first_top_level_function_or_method_ast @@ -115,15 +114,15 @@ def run(self) -> None: instrument_codeflash_trace_decorator(fto) trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" replay_tests_dir = Path(self.args.tests_root) / "codeflash_replay_tests" - trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Simply run all tests that use pytest-benchmark + trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Run all tests that use pytest-benchmark replay_count = generate_replay_test(trace_file, replay_tests_dir) if replay_count == 0: logger.info(f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization") else: function_benchmark_timings = get_function_benchmark_timings(trace_file) total_benchmark_timings = get_benchmark_timings(trace_file) - - print_benchmark_table(function_benchmark_timings, total_benchmark_timings) + function_to_results = validate_and_format_benchmark_table(function_benchmark_timings, total_benchmark_timings) + print_benchmark_table(function_to_results) logger.info("Finished tracing existing benchmarks") except Exception as e: logger.info(f"Error while tracing existing benchmarks: {e}") @@ -213,9 +212,12 @@ def run(self) -> None: f"Skipping optimization." ) continue - if self.args.benchmark and function_benchmark_timings and total_benchmark_timings: + qualified_name_w_module = function_to_optimize.qualified_name_with_modules_from_root( + self.args.project_root + ) + if self.args.benchmark and function_benchmark_timings and qualified_name_w_module in function_benchmark_timings and total_benchmark_timings: function_optimizer = self.create_function_optimizer( - function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code, function_benchmark_timings, total_benchmark_timings + function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code, function_benchmark_timings[qualified_name_w_module], total_benchmark_timings ) else: function_optimizer = self.create_function_optimizer( diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 248417452..888dca5c2 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -54,16 +54,18 @@ def to_console_string(self) -> str: continue total_benchmark_timing = self.total_benchmark_timings[benchmark_key] - # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values - benchmark_info += f"Original timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(total_benchmark_timing)}\n" - replay_speedup = self.replay_performance_gain - expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / ( - replay_speedup + 1) * og_benchmark_timing - benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(int(expected_new_benchmark_timing))}\n" - - benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing - benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 - benchmark_info += f"Benchmark speedup for {benchmark_file_name}::{benchmark_test_function}: {benchmark_speedup_percent:.2f}%\n\n" + if total_benchmark_timing == 0: + benchmark_info += f"Benchmark timing for {benchmark_file_name}::{benchmark_test_function} was improved, but the speedup cannot be estimated.\n" + else: + # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values + benchmark_info += f"Original timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(total_benchmark_timing)}\n" + replay_speedup = self.replay_performance_gain + expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / ( + replay_speedup + 1) * og_benchmark_timing + benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(int(expected_new_benchmark_timing))}\n" + benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing + benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 + benchmark_info += f"Benchmark speedup for {benchmark_file_name}::{benchmark_test_function}: {benchmark_speedup_percent:.2f}%\n\n" return ( f"Optimized {self.function_name} in {self.file_path}\n" diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 4113d954e..244b08029 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,9 +1,12 @@ import sqlite3 from codeflash.benchmarking.codeflash_trace import codeflash_trace +from codeflash.benchmarking.get_trace_info import get_function_benchmark_timings, get_benchmark_timings from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.replay_test import generate_replay_test from pathlib import Path + +from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table from codeflash.code_utils.code_utils import get_run_tmp_file import shutil @@ -11,7 +14,7 @@ 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" + benchmarks_root = project_root / "tests" / "pytest" / "benchmarks_test" tests_root = project_root / "tests" / "test_trace_benchmarks" tests_root.mkdir(parents=False, exist_ok=False) output_file = (tests_root / Path("test_trace_benchmarks.trace")).resolve() @@ -150,6 +153,62 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): """ assert test_sort_path.read_text("utf-8").strip()==test_sort_code.strip() + finally: + # cleanup + shutil.rmtree(tests_root) + pass + +def test_trace_multithreaded_benchmark() -> None: + project_root = Path(__file__).parent.parent / "code_to_optimize" + benchmarks_root = project_root / "tests" / "pytest" / "benchmarks_multithread" + tests_root = project_root / "tests" / "test_trace_benchmarks" + tests_root.mkdir(parents=False, exist_ok=False) + output_file = (tests_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_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM function_calls ORDER BY benchmark_file_name, benchmark_function_name, function_name") + function_calls = cursor.fetchall() + + # Assert the length of function calls + assert len(function_calls) == 10, f"Expected 10 function calls, but got {len(function_calls)}" + function_benchmark_timings = get_function_benchmark_timings(output_file) + total_benchmark_timings = get_benchmark_timings(output_file) + # This will throw an error if summed function timings exceed total benchmark timing + 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", "test_multithread_sort.py", 4), + ] + 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_name" + 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_file_name" + assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number" + # Close connection + conn.close() + finally: # cleanup shutil.rmtree(tests_root) From 6180c9de651576c9050864e4859992a9199b22e1 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 25 Mar 2025 09:51:12 -0700 Subject: [PATCH 025/122] refactored get_function_benchmark_timings and get_benchmark_timings into BenchmarkDatabaseUtils class --- .../benchmarking/benchmark_database_utils.py | 117 +++++++++++++++++ codeflash/benchmarking/get_trace_info.py | 121 ------------------ codeflash/optimization/optimizer.py | 6 +- tests/test_trace_benchmarks.py | 9 +- 4 files changed, 123 insertions(+), 130 deletions(-) delete mode 100644 codeflash/benchmarking/get_trace_info.py diff --git a/codeflash/benchmarking/benchmark_database_utils.py b/codeflash/benchmarking/benchmark_database_utils.py index b9b36079d..1c117553c 100644 --- a/codeflash/benchmarking/benchmark_database_utils.py +++ b/codeflash/benchmarking/benchmark_database_utils.py @@ -177,3 +177,120 @@ def close(self) -> None: self.connection.close() self.connection = None + + @staticmethod + def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, int]]: + """Process the trace file and extract timing data for all functions. + + Args: + trace_path: Path to the trace file + + Returns: + A nested dictionary where: + - Outer keys are module_name.qualified_name (module.class.function) + - Inner keys are benchmark filename :: benchmark test function :: line number + - Values are function timing in milliseconds + + """ + # Initialize the result dictionary + result = {} + + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + + try: + # Query the function_calls table for all function calls + cursor.execute( + "SELECT module_name, class_name, function_name, " + "benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns " + "FROM function_calls" + ) + + # Process each row + for row in cursor.fetchall(): + module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row + + # Create the function key (module_name.class_name.function_name) + if class_name: + qualified_name = f"{module_name}.{class_name}.{function_name}" + else: + qualified_name = f"{module_name}.{function_name}" + + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + + # Initialize the inner dictionary if needed + if qualified_name not in result: + result[qualified_name] = {} + + # If multiple calls to the same function in the same benchmark, + # add the times together + if benchmark_key in result[qualified_name]: + result[qualified_name][benchmark_key] += time_ns + else: + result[qualified_name][benchmark_key] = time_ns + + finally: + # Close the connection + connection.close() + + return result + + @staticmethod + def get_benchmark_timings(trace_path: Path) -> dict[str, int]: + """Extract total benchmark timings from trace files. + + Args: + trace_path: Path to the trace file + + Returns: + A dictionary mapping where: + - Keys are benchmark filename :: benchmark test function :: line number + - Values are total benchmark timing in milliseconds (with overhead subtracted) + + """ + # Initialize the result dictionary + result = {} + overhead_by_benchmark = {} + + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + + try: + # Query the function_calls table to get total overhead for each benchmark + cursor.execute( + "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " + "FROM function_calls " + "GROUP BY benchmark_file_name, 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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + 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_name, benchmark_function_name, benchmark_line_number, time_ns " + "FROM benchmark_timings" + ) + + # Process each row and subtract overhead + for row in cursor.fetchall(): + benchmark_file, benchmark_func, benchmark_line, time_ns = row + + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + + # Subtract overhead from total time + overhead = overhead_by_benchmark.get(benchmark_key, 0) + result[benchmark_key] = time_ns - overhead + + finally: + # Close the connection + connection.close() + + return result diff --git a/codeflash/benchmarking/get_trace_info.py b/codeflash/benchmarking/get_trace_info.py deleted file mode 100644 index d43327af7..000000000 --- a/codeflash/benchmarking/get_trace_info.py +++ /dev/null @@ -1,121 +0,0 @@ -import sqlite3 -from pathlib import Path - -from codeflash.discovery.functions_to_optimize import FunctionToOptimize - - -def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, int]]: - """Process the trace file and extract timing data for all functions. - - Args: - trace_path: Path to the trace file - - Returns: - A nested dictionary where: - - Outer keys are module_name.qualified_name (module.class.function) - - Inner keys are benchmark filename :: benchmark test function :: line number - - Values are function timing in milliseconds - - """ - # Initialize the result dictionary - result = {} - - # Connect to the SQLite database - connection = sqlite3.connect(trace_path) - cursor = connection.cursor() - - try: - # Query the function_calls table for all function calls - cursor.execute( - "SELECT module_name, class_name, function_name, " - "benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns " - "FROM function_calls" - ) - - # Process each row - for row in cursor.fetchall(): - module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row - - # Create the function key (module_name.class_name.function_name) - if class_name: - qualified_name = f"{module_name}.{class_name}.{function_name}" - else: - qualified_name = f"{module_name}.{function_name}" - - # Create the benchmark key (file::function::line) - benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - - # Initialize the inner dictionary if needed - if qualified_name not in result: - result[qualified_name] = {} - - # If multiple calls to the same function in the same benchmark, - # add the times together - if benchmark_key in result[qualified_name]: - result[qualified_name][benchmark_key] += time_ns - else: - result[qualified_name][benchmark_key] = time_ns - - finally: - # Close the connection - connection.close() - - return result - - -def get_benchmark_timings(trace_path: Path) -> dict[str, int]: - """Extract total benchmark timings from trace files. - - Args: - trace_path: Path to the trace file - - Returns: - A dictionary mapping where: - - Keys are benchmark filename :: benchmark test function :: line number - - Values are total benchmark timing in milliseconds (with overhead subtracted) - - """ - # Initialize the result dictionary - result = {} - overhead_by_benchmark = {} - - # Connect to the SQLite database - connection = sqlite3.connect(trace_path) - cursor = connection.cursor() - - try: - # Query the function_calls table to get total overhead for each benchmark - cursor.execute( - "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " - "FROM function_calls " - "GROUP BY benchmark_file_name, 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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - 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_name, benchmark_function_name, benchmark_line_number, time_ns " - "FROM benchmark_timings" - ) - - # Process each row and subtract overhead - for row in cursor.fetchall(): - benchmark_file, benchmark_func, benchmark_line, time_ns = row - - # Create the benchmark key (file::function::line) - benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - - # Subtract overhead from total time - overhead = overhead_by_benchmark.get(benchmark_key, 0) - result[benchmark_key] = time_ns - overhead - - finally: - # Close the connection - connection.close() - - return result diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 4a43a4c0b..8be62d963 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient +from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table @@ -24,7 +25,6 @@ from codeflash.telemetry.posthog_cf import ph from codeflash.verification.test_results import TestType 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 codeflash.benchmarking.instrument_codeflash_trace import instrument_codeflash_trace_decorator @@ -119,8 +119,8 @@ def run(self) -> None: if replay_count == 0: logger.info(f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization") else: - function_benchmark_timings = get_function_benchmark_timings(trace_file) - total_benchmark_timings = get_benchmark_timings(trace_file) + function_benchmark_timings = BenchmarkDatabaseUtils.get_function_benchmark_timings(trace_file) + total_benchmark_timings = BenchmarkDatabaseUtils.get_benchmark_timings(trace_file) function_to_results = validate_and_format_benchmark_table(function_benchmark_timings, total_benchmark_timings) print_benchmark_table(function_to_results) logger.info("Finished tracing existing benchmarks") diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 244b08029..fcc5b0f67 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,13 +1,11 @@ import sqlite3 -from codeflash.benchmarking.codeflash_trace import codeflash_trace -from codeflash.benchmarking.get_trace_info import get_function_benchmark_timings, get_benchmark_timings +from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.replay_test import generate_replay_test from pathlib import Path from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table -from codeflash.code_utils.code_utils import get_run_tmp_file import shutil @@ -180,9 +178,8 @@ def test_trace_multithreaded_benchmark() -> None: # Assert the length of function calls assert len(function_calls) == 10, f"Expected 10 function calls, but got {len(function_calls)}" - function_benchmark_timings = get_function_benchmark_timings(output_file) - total_benchmark_timings = get_benchmark_timings(output_file) - # This will throw an error if summed function timings exceed total benchmark timing + function_benchmark_timings = BenchmarkDatabaseUtils.get_function_benchmark_timings(output_file) + total_benchmark_timings = BenchmarkDatabaseUtils.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 From 67d3f1983ac929864fea4f2ef486a4dd8c88c6a2 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 25 Mar 2025 10:28:09 -0700 Subject: [PATCH 026/122] fixed isort --- codeflash/benchmarking/instrument_codeflash_trace.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py index 93f51baed..017cecaca 100644 --- a/codeflash/benchmarking/instrument_codeflash_trace.py +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -98,12 +98,12 @@ def instrument_codeflash_trace_decorator( """Instrument __init__ function with codeflash_trace decorator if it's in a class.""" # Instrument fto class original_code = function_to_optimize.file_path.read_text(encoding="utf-8") - - # Modify the code - modified_code = isort.code(add_codeflash_decorator_to_code( + new_code = add_codeflash_decorator_to_code( original_code, function_to_optimize - )) + ) + # Modify the code + modified_code = isort.code(code=new_code, float_to_top=True) # Write the modified code back to the file function_to_optimize.file_path.write_text(modified_code, encoding="utf-8") From f4be9becb51e997e801f7d85eba95f8c78835f5f Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 26 Mar 2025 10:20:28 -0700 Subject: [PATCH 027/122] modified PR info --- codeflash/benchmarking/utils.py | 61 +++++++++++++++++++- codeflash/github/PrComment.py | 15 ++++- codeflash/models/models.py | 41 +++++++++++++ codeflash/optimization/function_optimizer.py | 12 +++- codeflash/result/create_pr.py | 2 + codeflash/result/explanation.py | 33 +++-------- 6 files changed, 132 insertions(+), 32 deletions(-) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index eeacb6975..38c31b55b 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,7 +1,12 @@ +from __future__ import annotations +from typing import Optional + from rich.console import Console from rich.table import Table from codeflash.cli_cmds.console import logger +from codeflash.code_utils.time_utils import humanize_runtime +from codeflash.models.models import ProcessedBenchmarkInfo, BenchmarkDetail def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]], @@ -61,4 +66,58 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[str, float, ) # Print the table - console.print(table) \ No newline at end of file + console.print(table) + + +def process_benchmark_data( + replay_performance_gain: float, + fto_benchmark_timings: dict[str, int], + total_benchmark_timings: dict[str, int] +) -> Optional[ProcessedBenchmarkInfo]: + """Process benchmark data and generate detailed benchmark information. + + Args: + replay_performance_gain: The performance gain from replay + fto_benchmark_timings: Function to optimize benchmark timings + total_benchmark_timings: Total benchmark timings + + Returns: + ProcessedBenchmarkInfo containing processed benchmark details + + """ + if not replay_performance_gain or not fto_benchmark_timings or not total_benchmark_timings: + return None + + benchmark_details = [] + + for benchmark_key, og_benchmark_timing in fto_benchmark_timings.items(): + try: + benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") + except ValueError: + continue # Skip malformed benchmark keys + + total_benchmark_timing = total_benchmark_timings.get(benchmark_key, 0) + + if total_benchmark_timing == 0: + continue # Skip benchmarks with zero timing + + # Calculate expected new benchmark timing + expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + ( + 1 / (replay_performance_gain + 1) + ) * og_benchmark_timing + + # Calculate speedup + benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing + benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 + + benchmark_details.append( + BenchmarkDetail( + benchmark_name=benchmark_file_name, + test_function=benchmark_test_function, + original_timing=humanize_runtime(int(total_benchmark_timing)), + expected_new_timing=humanize_runtime(int(expected_new_benchmark_timing)), + speedup_percent=benchmark_speedup_percent + ) + ) + + return ProcessedBenchmarkInfo(benchmark_details=benchmark_details) \ No newline at end of file diff --git a/codeflash/github/PrComment.py b/codeflash/github/PrComment.py index f266a039d..4ef162cda 100644 --- a/codeflash/github/PrComment.py +++ b/codeflash/github/PrComment.py @@ -1,9 +1,11 @@ -from typing import Union +from __future__ import annotations +from typing import Union, Optional from pydantic import BaseModel from pydantic.dataclasses import dataclass from codeflash.code_utils.time_utils import humanize_runtime +from codeflash.models.models import BenchmarkDetail from codeflash.verification.test_results import TestResults @@ -18,15 +20,16 @@ class PrComment: speedup_pct: str winning_behavioral_test_results: TestResults winning_benchmarking_test_results: TestResults + benchmark_details: Optional[list[BenchmarkDetail]] = None - def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str]]: + def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str, Optional[list[dict[str, any]]]]]: report_table = { test_type.to_name(): result for test_type, result in self.winning_behavioral_test_results.get_test_pass_fail_report_by_type().items() if test_type.to_name() } - return { + result = { "optimization_explanation": self.optimization_explanation, "best_runtime": humanize_runtime(self.best_runtime), "original_runtime": humanize_runtime(self.original_runtime), @@ -38,6 +41,12 @@ def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str]]: "report_table": report_table, } + # Add benchmark details if available + if self.benchmark_details: + result["benchmark_details"] = self.benchmark_details + + return result + class FileDiffContent(BaseModel): oldContent: str diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 40e1f3906..52d1e4285 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -23,6 +23,7 @@ generate_candidates, ) from codeflash.code_utils.env_utils import is_end_to_end +from codeflash.code_utils.time_utils import humanize_runtime from codeflash.verification.test_results import TestResults, TestType # If the method spam is in the class Ham, which is at the top level of the module eggs in the package foo, the fully @@ -77,7 +78,47 @@ class BestOptimization(BaseModel): winning_benchmarking_test_results: TestResults winning_replay_benchmarking_test_results : Optional[TestResults] = None +@dataclass +class BenchmarkDetail: + benchmark_name: str + test_function: str + original_timing: str + expected_new_timing: str + speedup_percent: float + + def to_string(self) -> str: + return ( + f"Original timing for {self.benchmark_name}::{self.test_function}: {self.original_timing}\n" + f"Expected new timing for {self.benchmark_name}::{self.test_function}: {self.expected_new_timing}\n" + f"Benchmark speedup for {self.benchmark_name}::{self.test_function}: {self.speedup_percent:.2f}%\n" + ) + + def to_dict(self) -> dict[str, any]: + return { + "benchmark_name": self.benchmark_name, + "test_function": self.test_function, + "original_timing": self.original_timing, + "expected_new_timing": self.expected_new_timing, + "speedup_percent": self.speedup_percent + } +@dataclass +class ProcessedBenchmarkInfo: + benchmark_details: list[BenchmarkDetail] + + def to_string(self) -> str: + if not self.benchmark_details: + return "" + + result = "Benchmark Performance Details:\n" + for detail in self.benchmark_details: + result += detail.to_string() + "\n" + return result + + def to_dict(self) -> dict[str, list[dict[str, any]]]: + return { + "benchmark_details": [detail.to_dict() for detail in self.benchmark_details] + } class CodeString(BaseModel): code: Annotated[str, AfterValidator(validate_python_code)] file_path: Optional[Path] = None diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 70322c9fe..7474f6991 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -19,6 +19,7 @@ from rich.tree import Tree from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient +from codeflash.benchmarking.utils import process_benchmark_data from codeflash.cli_cmds.console import code_print, console, logger, progress_bar from codeflash.code_utils import env_utils from codeflash.code_utils.code_replacer import replace_function_definitions_in_module @@ -263,6 +264,13 @@ def optimize_function(self) -> Result[BestOptimization, str]: best_optimization.candidate.explanation, title="Best Candidate Explanation", border_style="blue" ) ) + processed_benchmark_info = None + if self.args.benchmark: + processed_benchmark_info = process_benchmark_data( + replay_performance_gain=best_optimization.replay_performance_gain, + fto_benchmark_timings=self.function_benchmark_timings, + total_benchmark_timings=self.total_benchmark_timings + ) explanation = Explanation( raw_explanation_message=best_optimization.candidate.explanation, winning_behavioral_test_results=best_optimization.winning_behavioral_test_results, @@ -271,9 +279,7 @@ def optimize_function(self) -> Result[BestOptimization, str]: best_runtime_ns=best_optimization.runtime, function_name=function_to_optimize_qualified_name, file_path=self.function_to_optimize.file_path, - replay_performance_gain=best_optimization.replay_performance_gain if self.args.benchmark else None, - fto_benchmark_timings = self.function_benchmark_timings if self.args.benchmark else None, - total_benchmark_timings = self.total_benchmark_timings if self.args.benchmark else None, + benchmark_details=processed_benchmark_info.benchmark_details if processed_benchmark_info else None ) self.log_successful_optimization(explanation, generated_tests) diff --git a/codeflash/result/create_pr.py b/codeflash/result/create_pr.py index e2d4da13c..da0c61961 100644 --- a/codeflash/result/create_pr.py +++ b/codeflash/result/create_pr.py @@ -77,6 +77,7 @@ def check_create_pr( speedup_pct=explanation.speedup_pct, winning_behavioral_test_results=explanation.winning_behavioral_test_results, winning_benchmarking_test_results=explanation.winning_benchmarking_test_results, + benchmark_details=explanation.benchmark_details ), existing_tests=existing_tests_source, generated_tests=generated_original_test_source, @@ -123,6 +124,7 @@ def check_create_pr( speedup_pct=explanation.speedup_pct, winning_behavioral_test_results=explanation.winning_behavioral_test_results, winning_benchmarking_test_results=explanation.winning_benchmarking_test_results, + benchmark_details=explanation.benchmark_details ), existing_tests=existing_tests_source, generated_tests=generated_original_test_source, diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 888dca5c2..10794991a 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -5,6 +5,7 @@ from pydantic.dataclasses import dataclass from codeflash.code_utils.time_utils import humanize_runtime +from codeflash.models.models import BenchmarkDetail from codeflash.verification.test_results import TestResults @@ -17,9 +18,7 @@ class Explanation: best_runtime_ns: int function_name: str file_path: Path - replay_performance_gain: Optional[float] - fto_benchmark_timings: Optional[dict[str, int]] - total_benchmark_timings: Optional[dict[str, int]] + benchmark_details: Optional[list[BenchmarkDetail]] = None @property def perf_improvement_line(self) -> str: @@ -43,29 +42,13 @@ def to_console_string(self) -> str: original_runtime_human = humanize_runtime(self.original_runtime_ns) best_runtime_human = humanize_runtime(self.best_runtime_ns) benchmark_info = "" - if self.replay_performance_gain and self.fto_benchmark_timings and self.total_benchmark_timings: - benchmark_info += "Benchmark Performance Details:\n" - for benchmark_key, og_benchmark_timing in self.fto_benchmark_timings.items(): - # benchmark key is benchmark filename :: benchmark test function :: line number - try: - benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") - except ValueError: - benchmark_info += f"Benchmark key {benchmark_key} is not in the expected format.\n" - continue - total_benchmark_timing = self.total_benchmark_timings[benchmark_key] - if total_benchmark_timing == 0: - benchmark_info += f"Benchmark timing for {benchmark_file_name}::{benchmark_test_function} was improved, but the speedup cannot be estimated.\n" - else: - # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values - benchmark_info += f"Original timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(total_benchmark_timing)}\n" - replay_speedup = self.replay_performance_gain - expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / ( - replay_speedup + 1) * og_benchmark_timing - benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(int(expected_new_benchmark_timing))}\n" - benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing - benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 - benchmark_info += f"Benchmark speedup for {benchmark_file_name}::{benchmark_test_function}: {benchmark_speedup_percent:.2f}%\n\n" + if self.benchmark_details: + benchmark_info += "Benchmark Performance Details:\n" + for detail in self.benchmark_details: + benchmark_info += f"Original timing for {detail.benchmark_name}::{detail.test_function}: {detail.original_timing}\n" + benchmark_info += f"Expected new timing for {detail.benchmark_name}::{detail.test_function}: {detail.expected_new_timing}\n" + benchmark_info += f"Benchmark speedup for {detail.benchmark_name}::{detail.test_function}: {detail.speedup_percent:.2f}%\n\n" return ( f"Optimized {self.function_name} in {self.file_path}\n" From 77f43a58ac9d2dbee1f3e61c23656c013db806aa Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 26 Mar 2025 12:08:52 -0700 Subject: [PATCH 028/122] mypy fix --- codeflash/github/PrComment.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/codeflash/github/PrComment.py b/codeflash/github/PrComment.py index 4ef162cda..5b891b8a5 100644 --- a/codeflash/github/PrComment.py +++ b/codeflash/github/PrComment.py @@ -22,14 +22,14 @@ class PrComment: winning_benchmarking_test_results: TestResults benchmark_details: Optional[list[BenchmarkDetail]] = None - def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str, Optional[list[dict[str, any]]]]]: + def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str, Optional[list[BenchmarkDetail]]]]: report_table = { test_type.to_name(): result for test_type, result in self.winning_behavioral_test_results.get_test_pass_fail_report_by_type().items() if test_type.to_name() } - result = { + return { "optimization_explanation": self.optimization_explanation, "best_runtime": humanize_runtime(self.best_runtime), "original_runtime": humanize_runtime(self.original_runtime), @@ -39,14 +39,9 @@ def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str, Option "speedup_pct": self.speedup_pct, "loop_count": self.winning_benchmarking_test_results.number_of_loops(), "report_table": report_table, + "benchmark_details": self.benchmark_details if self.benchmark_details else None, } - # Add benchmark details if available - if self.benchmark_details: - result["benchmark_details"] = self.benchmark_details - - return result - class FileDiffContent(BaseModel): oldContent: str From da6385f1782887fc7199f57743c5b42eec0f67a3 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 26 Mar 2025 15:46:49 -0700 Subject: [PATCH 029/122] use dill instead of pickle --- codeflash/benchmarking/codeflash_trace.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index f708d752f..3b55aa6ba 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -1,6 +1,6 @@ import functools import os -import pickle +import dill as pickle import time from typing import Callable @@ -63,13 +63,11 @@ def wrapper(*args, **kwargs): overhead_end_time = time.thread_time_ns() overhead_time = overhead_end_time - overhead_start_time - self.function_calls_data.append( (func.__name__, class_name, func.__module__, func.__code__.co_filename, benchmark_function_name, benchmark_file_name, benchmark_line_number, execution_time, overhead_time, pickled_args, pickled_kwargs) ) - print("appended") except Exception as e: print(f"Error in codeflash_trace: {e}") From f34f22fde8ddab7adfce9e218afed8175e324bf1 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 28 Mar 2025 13:55:59 -0700 Subject: [PATCH 030/122] modified the benchmarking approach. codeflash_trace and codeflash_benchmark_plugins are singleton instances that will both handle writing to disk. enables flushing to disk once a limit is reached. also added various details to the tracer --- .../benchmarking/benchmark_database_utils.py | 296 ------------------ codeflash/benchmarking/codeflash_trace.py | 167 +++++++--- .../instrument_codeflash_trace.py | 85 +++-- codeflash/benchmarking/plugin/plugin.py | 262 ++++++++++++++-- .../pytest_new_process_trace_benchmarks.py | 16 +- codeflash/benchmarking/replay_test.py | 12 +- codeflash/benchmarking/trace_benchmarks.py | 5 +- codeflash/benchmarking/utils.py | 20 +- codeflash/cli_cmds/cli.py | 1 - codeflash/discovery/functions_to_optimize.py | 1 - codeflash/models/models.py | 11 + codeflash/optimization/function_optimizer.py | 19 +- codeflash/optimization/optimizer.py | 16 +- codeflash/verification/test_results.py | 68 ++-- tests/test_instrument_codeflash_trace.py | 243 +++++++++++++- tests/test_trace_benchmarks.py | 14 +- 16 files changed, 718 insertions(+), 518 deletions(-) delete mode 100644 codeflash/benchmarking/benchmark_database_utils.py diff --git a/codeflash/benchmarking/benchmark_database_utils.py b/codeflash/benchmarking/benchmark_database_utils.py deleted file mode 100644 index 1c117553c..000000000 --- a/codeflash/benchmarking/benchmark_database_utils.py +++ /dev/null @@ -1,296 +0,0 @@ -import sqlite3 -from pathlib import Path - -import pickle - - -class BenchmarkDatabaseUtils: - def __init__(self, trace_path :Path) -> None: - self.trace_path = trace_path - self.connection = None - - def setup(self) -> None: - try: - # Open connection - self.connection = sqlite3.connect(self.trace_path) - cur = self.connection.cursor() - cur.execute("PRAGMA synchronous = OFF") - cur.execute( - "CREATE TABLE IF NOT EXISTS function_calls(" - "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," - "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," - "time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" - ) - cur.execute( - "CREATE TABLE IF NOT EXISTS benchmark_timings(" - "benchmark_file_name TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER," - "time_ns INTEGER)" # Added closing parenthesis - ) - self.connection.commit() - # Don't close the connection here - except Exception as e: - print(f"Database setup error: {e}") - if self.connection: - self.connection.close() - self.connection = None - raise - - def write_function_timings(self, data: list[tuple]) -> None: - if not self.connection: - self.connection = sqlite3.connect(self.trace_path) - - try: - cur = self.connection.cursor() - # Insert data into the function_calls table - cur.executemany( - "INSERT INTO function_calls " - "(function_name, class_name, module_name, file_name, benchmark_function_name, " - "benchmark_file_name, benchmark_line_number, time_ns, overhead_time_ns, args, kwargs) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - data - ) - self.connection.commit() - except Exception as e: - print(f"Error writing to function timings database: {e}") - self.connection.rollback() - raise - - def write_benchmark_timings(self, data: list[tuple]) -> None: - if not self.connection: - self.connection = sqlite3.connect(self.trace_path) - - try: - cur = self.connection.cursor() - # Insert data into the benchmark_timings table - cur.executemany( - "INSERT INTO benchmark_timings (benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns) VALUES (?, ?, ?, ?)", - data - ) - self.connection.commit() - except Exception as e: - print(f"Error writing to benchmark timings database: {e}") - self.connection.rollback() - raise - - def print_function_timings(self, limit: int = None) -> None: - """Print the contents of a CodeflashTrace SQLite database. - - Args: - limit: Maximum number of records to print (None for all) - """ - if not self.connection: - self.connection = sqlite3.connect(self.trace_path) - try: - cur = self.connection.cursor() - - # Get the count of records - cur.execute("SELECT COUNT(*) FROM function_calls") - total_records = cur.fetchone()[0] - print(f"Found {total_records} function call records in {self.trace_path}") - - # Build the query with optional limit - query = "SELECT * FROM function_calls" - if limit: - query += f" LIMIT {limit}" - - # Execute the query - cur.execute(query) - - # Print column names - columns = [desc[0] for desc in cur.description] - print("\nColumns:", columns) - print("\n" + "=" * 80 + "\n") - - # Print each row - for i, row in enumerate(cur.fetchall()): - print(f"Record #{i + 1}:") - print(f" Function: {row[0]}") - print(f" Class: {row[1]}") - print(f" Module: {row[2]}") - print(f" File: {row[3]}") - print(f" Benchmark Function: {row[4] or 'N/A'}") - print(f" Benchmark File: {row[5] or 'N/A'}") - print(f" Benchmark Line: {row[6] or 'N/A'}") - print(f" Execution Time: {row[7]:.6f} seconds") - print(f" Overhead Time: {row[8]:.6f} seconds") - - # Unpickle and print args and kwargs - try: - args = pickle.loads(row[9]) - kwargs = pickle.loads(row[10]) - - print(f" Args: {args}") - print(f" Kwargs: {kwargs}") - except Exception as e: - print(f" Error unpickling args/kwargs: {e}") - print(f" Raw args: {row[9]}") - print(f" Raw kwargs: {row[10]}") - - print("\n" + "-" * 40 + "\n") - - except Exception as e: - print(f"Error reading database: {e}") - - def print_benchmark_timings(self, limit: int = None) -> None: - """Print the contents of a CodeflashTrace SQLite database. - Args: - limit: Maximum number of records to print (None for all) - """ - if not self.connection: - self.connection = sqlite3.connect(self.trace_path) - try: - cur = self.connection.cursor() - - # Get the count of records - cur.execute("SELECT COUNT(*) FROM benchmark_timings") - total_records = cur.fetchone()[0] - print(f"Found {total_records} benchmark timing records in {self.trace_path}") - - # Build the query with optional limit - query = "SELECT * FROM benchmark_timings" - if limit: - query += f" LIMIT {limit}" - - # Execute the query - cur.execute(query) - - # Print column names - columns = [desc[0] for desc in cur.description] - print("\nColumns:", columns) - print("\n" + "=" * 80 + "\n") - - # Print each row - for i, row in enumerate(cur.fetchall()): - print(f"Record #{i + 1}:") - print(f" Benchmark File: {row[0] or 'N/A'}") - print(f" Benchmark Function: {row[1] or 'N/A'}") - print(f" Benchmark Line: {row[2] or 'N/A'}") - print(f" Execution Time: {row[3] / 1e9:.6f} seconds") # Convert nanoseconds to seconds - print("\n" + "-" * 40 + "\n") - - except Exception as e: - print(f"Error reading benchmark timings database: {e}") - - - def close(self) -> None: - if self.connection: - self.connection.close() - self.connection = None - - - @staticmethod - def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, int]]: - """Process the trace file and extract timing data for all functions. - - Args: - trace_path: Path to the trace file - - Returns: - A nested dictionary where: - - Outer keys are module_name.qualified_name (module.class.function) - - Inner keys are benchmark filename :: benchmark test function :: line number - - Values are function timing in milliseconds - - """ - # Initialize the result dictionary - result = {} - - # Connect to the SQLite database - connection = sqlite3.connect(trace_path) - cursor = connection.cursor() - - try: - # Query the function_calls table for all function calls - cursor.execute( - "SELECT module_name, class_name, function_name, " - "benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns " - "FROM function_calls" - ) - - # Process each row - for row in cursor.fetchall(): - module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row - - # Create the function key (module_name.class_name.function_name) - if class_name: - qualified_name = f"{module_name}.{class_name}.{function_name}" - else: - qualified_name = f"{module_name}.{function_name}" - - # Create the benchmark key (file::function::line) - benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - - # Initialize the inner dictionary if needed - if qualified_name not in result: - result[qualified_name] = {} - - # If multiple calls to the same function in the same benchmark, - # add the times together - if benchmark_key in result[qualified_name]: - result[qualified_name][benchmark_key] += time_ns - else: - result[qualified_name][benchmark_key] = time_ns - - finally: - # Close the connection - connection.close() - - return result - - @staticmethod - def get_benchmark_timings(trace_path: Path) -> dict[str, int]: - """Extract total benchmark timings from trace files. - - Args: - trace_path: Path to the trace file - - Returns: - A dictionary mapping where: - - Keys are benchmark filename :: benchmark test function :: line number - - Values are total benchmark timing in milliseconds (with overhead subtracted) - - """ - # Initialize the result dictionary - result = {} - overhead_by_benchmark = {} - - # Connect to the SQLite database - connection = sqlite3.connect(trace_path) - cursor = connection.cursor() - - try: - # Query the function_calls table to get total overhead for each benchmark - cursor.execute( - "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " - "FROM function_calls " - "GROUP BY benchmark_file_name, 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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - 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_name, benchmark_function_name, benchmark_line_number, time_ns " - "FROM benchmark_timings" - ) - - # Process each row and subtract overhead - for row in cursor.fetchall(): - benchmark_file, benchmark_func, benchmark_line, time_ns = row - - # Create the benchmark key (file::function::line) - benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - - # Subtract overhead from total time - overhead = overhead_by_benchmark.get(benchmark_key, 0) - result[benchmark_key] = time_ns - overhead - - finally: - # Close the connection - connection.close() - - return result diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 3b55aa6ba..2ae57307b 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -1,23 +1,90 @@ import functools import os -import dill as pickle -import time -from typing import Callable - +import sqlite3 +import sys +import pickle +import dill +import time +from typing import Callable, Optional class CodeflashTrace: - """A class that provides both a decorator for tracing function calls - and a context manager for managing the tracing data lifecycle. - """ + """Decorator class that traces and profiles function execution.""" def __init__(self) -> None: self.function_calls_data = [] + self.function_call_count = 0 + self.pickle_count_limit = 1000 + self._connection = None + self._trace_path = None + + def setup(self, trace_path: str) -> None: + """Set up the database connection for direct writing. + + Args: + trace_path: Path to the trace database file - def __exit__(self, exc_type, exc_val, exc_tb) -> None: - # Cleanup is optional here - pass + """ + try: + self._trace_path = trace_path + self._connection = sqlite3.connect(self._trace_path) + cur = self._connection.cursor() + cur.execute("PRAGMA synchronous = OFF") + cur.execute( + "CREATE TABLE IF NOT EXISTS benchmark_function_timings(" + "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," + "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," + "function_time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" + ) + self._connection.commit() + except Exception as e: + print(f"Database setup error: {e}") + if self._connection: + self._connection.close() + self._connection = None + raise + + def write_function_timings(self) -> None: + """Write function call data directly to the database. + + Args: + data: List of function call data tuples to write + """ + if not self.function_calls_data: + return # No data to write + + if self._connection is None and self._trace_path is not None: + self._connection = sqlite3.connect(self._trace_path) + + try: + cur = self._connection.cursor() + # Insert data into the benchmark_function_timings table + cur.executemany( + "INSERT INTO benchmark_function_timings" + "(function_name, class_name, module_name, file_name, benchmark_function_name, " + "benchmark_file_name, benchmark_line_number, function_time_ns, overhead_time_ns, args, kwargs) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + self.function_calls_data + ) + self._connection.commit() + self.function_calls_data = [] + except Exception as e: + print(f"Error writing to function timings database: {e}") + if self._connection: + self._connection.rollback() + raise + + def open(self) -> None: + """Open the database connection.""" + if self._connection is None: + self._connection = sqlite3.connect(self._trace_path) + + def close(self) -> None: + """Close the database connection.""" + if self._connection: + self._connection.close() + self._connection = None def __call__(self, func: Callable) -> Callable: """Use as a decorator to trace function execution. @@ -38,39 +105,55 @@ def wrapper(*args, **kwargs): # Calculate execution time execution_time = end_time - start_time - # Measure overhead - overhead_start_time = time.thread_time_ns() - - try: - # Check if currently in pytest benchmark fixture - if os.environ.get("CODEFLASH_BENCHMARKING", "False") == "False": - return result - - # 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", "") - benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "") - # Get class name - class_name = "" - qualname = func.__qualname__ - if "." in qualname: - class_name = qualname.split(".")[0] - # Calculate overhead time - overhead_end_time = time.thread_time_ns() - overhead_time = overhead_end_time - overhead_start_time - - self.function_calls_data.append( - (func.__name__, class_name, func.__module__, func.__code__.co_filename, - benchmark_function_name, benchmark_file_name, benchmark_line_number, execution_time, - overhead_time, pickled_args, pickled_kwargs) - ) - except Exception as e: - print(f"Error in codeflash_trace: {e}") + 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": + return result + + # 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", "") + benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "") + # Get class name + class_name = "" + qualname = func.__qualname__ + if "." in qualname: + class_name = qualname.split(".")[0] + + if self.function_call_count <= self.pickle_count_limit: + 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) + 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 + + if len(self.function_calls_data) > 1000: + self.write_function_timings() + # Calculate overhead time + 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_file_name, benchmark_line_number, execution_time, + overhead_time, pickled_args, pickled_kwargs) + ) return result return wrapper diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py index 017cecaca..06e93daf8 100644 --- a/codeflash/benchmarking/instrument_codeflash_trace.py +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -1,3 +1,5 @@ +from pathlib import Path + import isort import libcst as cst @@ -5,40 +7,35 @@ class AddDecoratorTransformer(cst.CSTTransformer): - def __init__(self, function_name, class_name=None): + def __init__(self, target_functions: set[tuple[str, str]]) -> None: super().__init__() - self.function_name = function_name - self.class_name = class_name - self.in_target_class = (class_name is None) # If no class name, always "in target class" + self.target_functions = target_functions self.added_codeflash_trace = False + self.class_name = "" + self.decorator = cst.Decorator( + decorator=cst.Name(value="codeflash_trace") + ) def leave_ClassDef(self, original_node, updated_node): - if self.class_name and original_node.name.value == self.class_name: - self.in_target_class = False + self.class_name = "" return updated_node def visit_ClassDef(self, node): - if self.class_name and node.name.value == self.class_name: - self.in_target_class = True - return True + if self.class_name: # Don't go into nested class + return False + self.class_name = node.name.value def leave_FunctionDef(self, original_node, updated_node): - if not self.in_target_class or original_node.name.value != self.function_name: + if (self.class_name, original_node.name.value) in self.target_functions: + # Add the new decorator after any existing decorators, so it gets executed first + updated_decorators = list(updated_node.decorators) + [self.decorator] + self.added_codeflash_trace = True + return updated_node.with_changes( + decorators=updated_decorators + ) + else: return updated_node - # Create the codeflash_trace decorator - decorator = cst.Decorator( - decorator=cst.Name(value="codeflash_trace") - ) - - # Add the new decorator after any existing decorators - updated_decorators = list(updated_node.decorators) + [decorator] - self.added_codeflash_trace = True - # Return the updated node with the new decorator - return updated_node.with_changes( - decorators=updated_decorators - ) - def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: # Create import statement for codeflash_trace if not self.added_codeflash_trace: @@ -62,12 +59,12 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c ] ) - # Insert at the beginning of the file + # Insert at the beginning of the file. We'll use isort later to sort the imports. new_body = [import_stmt, *list(updated_node.body)] return updated_node.with_changes(body=new_body) -def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToOptimize) -> str: +def add_codeflash_decorator_to_code(code: str, functions_to_optimize: list[FunctionToOptimize]) -> str: """Add codeflash_trace to a function. Args: @@ -76,15 +73,17 @@ def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToO Returns: The modified source code as a string + """ - # Extract class name if present - class_name = None - if len(function_to_optimize.parents) == 1 and function_to_optimize.parents[0].type == "ClassDef": - class_name = function_to_optimize.parents[0].name + target_functions = set() + for function_to_optimize in functions_to_optimize: + class_name = "" + if len(function_to_optimize.parents) == 1 and function_to_optimize.parents[0].type == "ClassDef": + class_name = function_to_optimize.parents[0].name + target_functions.add((class_name, function_to_optimize.function_name)) transformer = AddDecoratorTransformer( - function_name=function_to_optimize.function_name, - class_name=class_name + target_functions = target_functions, ) module = cst.parse_module(code) @@ -93,17 +92,17 @@ def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToO def instrument_codeflash_trace_decorator( - function_to_optimize: FunctionToOptimize + file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]] ) -> None: - """Instrument __init__ function with codeflash_trace decorator if it's in a class.""" - # Instrument fto class - original_code = function_to_optimize.file_path.read_text(encoding="utf-8") - new_code = add_codeflash_decorator_to_code( - original_code, - function_to_optimize - ) - # Modify the code - modified_code = isort.code(code=new_code, float_to_top=True) + """Instrument codeflash_trace decorator to functions to optimize.""" + for file_path, functions_to_optimize in file_to_funcs_to_optimize.items(): + original_code = file_path.read_text(encoding="utf-8") + new_code = add_codeflash_decorator_to_code( + original_code, + functions_to_optimize + ) + # Modify the code + modified_code = isort.code(code=new_code, float_to_top=True) - # Write the modified code back to the file - function_to_optimize.file_path.write_text(modified_code, encoding="utf-8") + # Write the modified code back to the file + file_path.write_text(modified_code, encoding="utf-8") diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index ee7504ec4..9d7da6ef2 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -1,33 +1,195 @@ +from __future__ import annotations +import os +import sqlite3 import sys - -import pytest import time -import os +from pathlib import Path +import pytest +from codeflash.benchmarking.codeflash_trace import codeflash_trace +from codeflash.models.models import BenchmarkKey + + class CodeFlashBenchmarkPlugin: - benchmark_timings = [] + def __init__(self) -> None: + self._trace_path = None + self._connection = None + self.benchmark_timings = [] - class Benchmark: - def __init__(self, request): - self.request = request + def setup(self, trace_path:str) -> None: + try: + # Open connection + self._trace_path = trace_path + self._connection = sqlite3.connect(self._trace_path) + cur = self._connection.cursor() + cur.execute("PRAGMA synchronous = OFF") + cur.execute( + "CREATE TABLE IF NOT EXISTS benchmark_timings(" + "benchmark_file_name TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER," + "benchmark_time_ns INTEGER)" + ) + self._connection.commit() + self.close() # Reopen only at the end of pytest session + except Exception as e: + print(f"Database setup error: {e}") + if self._connection: + self._connection.close() + self._connection = None + raise - def __call__(self, func, *args, **kwargs): - benchmark_file_name = self.request.node.fspath.basename - benchmark_function_name = self.request.node.name - line_number = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack + def write_benchmark_timings(self) -> None: + if not self.benchmark_timings: + return # No data to write - os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name - os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = benchmark_file_name - os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = line_number - os.environ["CODEFLASH_BENCHMARKING"] = "True" + if self._connection is None: + self._connection = sqlite3.connect(self._trace_path) - start = time.perf_counter_ns() - result = func(*args, **kwargs) - end = time.perf_counter_ns() + try: + cur = self._connection.cursor() + # Insert data into the benchmark_timings table + cur.executemany( + "INSERT INTO benchmark_timings (benchmark_file_name, benchmark_function_name, benchmark_line_number, benchmark_time_ns) VALUES (?, ?, ?, ?)", + self.benchmark_timings + ) + self._connection.commit() + self.benchmark_timings = [] # Clear the benchmark timings list + except Exception as e: + print(f"Error writing to benchmark timings database: {e}") + self._connection.rollback() + raise + def close(self) -> None: + if self._connection: + self._connection.close() + self._connection = None + + @staticmethod + def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[BenchmarkKey, int]]: + """Process the trace file and extract timing data for all functions. + + Args: + trace_path: Path to the trace file + + Returns: + A nested dictionary where: + - Outer keys are module_name.qualified_name (module.class.function) + - Inner keys are of type BenchmarkKey + - Values are function timing in milliseconds + + """ + # Initialize the result dictionary + result = {} + + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + + try: + # Query the function_calls table for all function calls + cursor.execute( + "SELECT module_name, class_name, function_name, " + "benchmark_file_name, benchmark_function_name, benchmark_line_number, function_time_ns " + "FROM benchmark_function_timings" + ) + + # Process each row + for row in cursor.fetchall(): + module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row + + # Create the function key (module_name.class_name.function_name) + if class_name: + qualified_name = f"{module_name}.{class_name}.{function_name}" + else: + qualified_name = f"{module_name}.{function_name}" + + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + # Initialize the inner dictionary if needed + if qualified_name not in result: + result[qualified_name] = {} + + # If multiple calls to the same function in the same benchmark, + # add the times together + if benchmark_key in result[qualified_name]: + result[qualified_name][benchmark_key] += time_ns + else: + result[qualified_name][benchmark_key] = time_ns + + finally: + # Close the connection + connection.close() + + return result + + @staticmethod + def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]: + """Extract total benchmark timings from trace files. + + Args: + trace_path: Path to the trace file + + Returns: + A dictionary mapping where: + - Keys are of type BenchmarkKey + - Values are total benchmark timing in milliseconds (with overhead subtracted) + + """ + # Initialize the result dictionary + result = {} + overhead_by_benchmark = {} + + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + + try: + # Query the benchmark_function_timings table to get total overhead for each benchmark + cursor.execute( + "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " + "FROM benchmark_function_timings " + "GROUP BY benchmark_file_name, 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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + 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_name, benchmark_function_name, benchmark_line_number, benchmark_time_ns " + "FROM benchmark_timings" + ) + + # Process each row and subtract overhead + for row in cursor.fetchall(): + benchmark_file, benchmark_func, benchmark_line, time_ns = row + + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + # Subtract overhead from total time + overhead = overhead_by_benchmark.get(benchmark_key, 0) + result[benchmark_key] = time_ns - overhead + + finally: + # Close the connection + connection.close() + + return result + + # Pytest hooks + @pytest.hookimpl + def pytest_sessionfinish(self, session, exitstatus): + """Execute after whole test run is completed.""" + # Write any remaining benchmark timings to the database + codeflash_trace.close() + if self.benchmark_timings: + self.write_benchmark_timings() + # Close the database connection + self.close() - os.environ["CODEFLASH_BENCHMARKING"] = "False" - CodeFlashBenchmarkPlugin.benchmark_timings.append( - (benchmark_file_name, benchmark_function_name, line_number, end - start)) - return result @staticmethod def pytest_addoption(parser): parser.addoption( @@ -39,11 +201,13 @@ def pytest_addoption(parser): @staticmethod def pytest_plugin_registered(plugin, manager): + # Not necessary since run with -p no:benchmark, but just in case if hasattr(plugin, "name") and plugin.name == "pytest-benchmark": manager.unregister(plugin) @staticmethod def pytest_collection_modifyitems(config, items): + # Skip tests that don't have the benchmark fixture if not config.getoption("--codeflash-trace"): return @@ -53,10 +217,62 @@ def pytest_collection_modifyitems(config, items): continue 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 + + """ + benchmark_file_name = self.request.node.fspath.basename + 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_NAME"] = benchmark_file_name + os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(line_number) + os.environ["CODEFLASH_BENCHMARKING"] = "True" + + # Run the function + start = time.perf_counter_ns() + result = func(*args, **kwargs) + end = time.perf_counter_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 + codeflash_trace.function_call_count = 0 + # Add to the benchmark timings buffer + codeflash_benchmark_plugin.benchmark_timings.append( + (benchmark_file_name, benchmark_function_name, line_number, end - start)) + + return result + @staticmethod @pytest.fixture def benchmark(request): if not request.config.getoption("--codeflash-trace"): return None - return CodeFlashBenchmarkPlugin.Benchmark(request) \ No newline at end of file + return CodeFlashBenchmarkPlugin.Benchmark(request) + +codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin() \ No newline at end of file diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 6d4c85f41..7b6bd747a 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -1,11 +1,8 @@ import sys from pathlib import Path -from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils -from codeflash.verification.verification_utils import get_test_file_path -from plugin.plugin import CodeFlashBenchmarkPlugin from codeflash.benchmarking.codeflash_trace import codeflash_trace -from codeflash.code_utils.code_utils import get_run_tmp_file +from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin benchmarks_root = sys.argv[1] tests_root = sys.argv[2] @@ -16,16 +13,11 @@ import pytest try: - db = BenchmarkDatabaseUtils(trace_path=Path(trace_file)) - db.setup() + codeflash_benchmark_plugin.setup(trace_file) + codeflash_trace.setup(trace_file) exitcode = pytest.main( - [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashBenchmarkPlugin()] + [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[codeflash_benchmark_plugin] ) # Errors will be printed to stdout, not stderr - db.write_function_timings(codeflash_trace.function_calls_data) - db.write_benchmark_timings(CodeFlashBenchmarkPlugin.benchmark_timings) - # db.print_function_timings() - # db.print_benchmark_timings() - db.close() except Exception as e: print(f"Failed to collect tests: {e!s}", file=sys.stderr) diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index a1d5b370a..670d6e4bd 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -13,7 +13,7 @@ from pathlib import Path def get_next_arg_and_return( - trace_file: str, function_name: str, file_name: str, class_name: str | None = None, num_to_get: int = 25 + trace_file: str, function_name: str, file_name: str, class_name: str | None = None, num_to_get: int = 256 ) -> Generator[Any]: db = sqlite3.connect(trace_file) cur = db.cursor() @@ -21,12 +21,12 @@ def get_next_arg_and_return( if class_name is not None: cursor = cur.execute( - "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? AND class_name = ? ORDER BY time_ns ASC LIMIT ?", + "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_name = ? AND class_name = ? LIMIT ?", (function_name, file_name, class_name, limit), ) else: cursor = cur.execute( - "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? AND class_name = '' ORDER BY time_ns ASC LIMIT ?", + "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_name = ? AND class_name = '' LIMIT ?", (function_name, file_name, limit), ) @@ -42,7 +42,7 @@ def create_trace_replay_test_code( trace_file: str, functions_data: list[dict[str, Any]], test_framework: str = "pytest", - max_run_count=100 + max_run_count=256 ) -> str: """Create a replay test for functions based on trace data. @@ -217,7 +217,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework # Get distinct benchmark names cursor.execute( - "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM function_calls" + "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM benchmark_function_timings" ) benchmarks = cursor.fetchall() @@ -226,7 +226,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework benchmark_function_name, benchmark_file_name = benchmark # Get functions associated with this benchmark cursor.execute( - "SELECT DISTINCT function_name, class_name, module_name, file_name, benchmark_line_number FROM function_calls " + "SELECT DISTINCT function_name, class_name, module_name, file_name, benchmark_line_number FROM benchmark_function_timings " "WHERE benchmark_function_name = ? AND benchmark_file_name = ?", (benchmark_function_name, benchmark_file_name) ) diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 79395db79..8882078d9 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -9,7 +9,7 @@ from pathlib import Path import subprocess -def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, trace_file: Path) -> None: +def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, trace_file: Path, timeout:int = 300) -> None: result = subprocess.run( [ SAFE_SYS_EXECUTABLE, @@ -23,6 +23,7 @@ def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root capture_output=True, text=True, env={"PYTHONPATH": str(project_root)}, + timeout=timeout, ) if result.returncode != 0: if "ERROR collecting" in result.stdout: @@ -38,5 +39,5 @@ def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root else: error_section = result.stdout logger.warning( - f"Error collecting benchmarks - Pytest Exit code: {result.returncode}={ExitCode(result.returncode).name}\n {error_section}" + f"Error collecting benchmarks - Pytest Exit code: {result.returncode}, {error_section}" ) \ No newline at end of file diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 38c31b55b..5f14f141f 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -6,29 +6,30 @@ from codeflash.cli_cmds.console import logger from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.models.models import ProcessedBenchmarkInfo, BenchmarkDetail +from codeflash.models.models import ProcessedBenchmarkInfo, BenchmarkDetail, BenchmarkKey +from codeflash.result.critic import performance_gain -def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]], - total_benchmark_timings: dict[str, int]) -> dict[str, list[tuple[str, float, float, float]]]: +def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[BenchmarkKey, int]], + total_benchmark_timings: dict[BenchmarkKey, int]) -> dict[str, list[tuple[str, float, float, float]]]: function_to_result = {} # Process each function's benchmark data for func_path, test_times in function_benchmark_timings.items(): # Sort by percentage (highest first) sorted_tests = [] - for test_name, func_time in test_times.items(): - total_time = total_benchmark_timings.get(test_name, 0) + for benchmark_key, func_time in test_times.items(): + total_time = total_benchmark_timings.get(benchmark_key, 0) if func_time > total_time: - logger.debug(f"Skipping test {test_name} due to func_time {func_time} > total_time {total_time}") + logger.debug(f"Skipping test {benchmark_key} due to func_time {func_time} > total_time {total_time}") # If the function time is greater than total time, likely to have multithreading / multiprocessing issues. # Do not try to project the optimization impact for this function. - sorted_tests.append((test_name, 0.0, 0.0, 0.0)) + sorted_tests.append((str(benchmark_key), 0.0, 0.0, 0.0)) if total_time > 0: percentage = (func_time / total_time) * 100 # Convert nanoseconds to milliseconds func_time_ms = func_time / 1_000_000 total_time_ms = total_time / 1_000_000 - sorted_tests.append((test_name, total_time_ms, func_time_ms, percentage)) + sorted_tests.append((str(benchmark_key), total_time_ms, func_time_ms, percentage)) sorted_tests.sort(key=lambda x: x[3], reverse=True) function_to_result[func_path] = sorted_tests return function_to_result @@ -107,8 +108,7 @@ def process_benchmark_data( ) * og_benchmark_timing # Calculate speedup - benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing - benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 + benchmark_speedup_percent = performance_gain(original_runtime_ns=total_benchmark_timing, optimized_runtime_ns=int(expected_new_benchmark_timing)) * 100 benchmark_details.append( BenchmarkDetail( diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 96bb0cef3..d1e786703 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -121,7 +121,6 @@ def process_pyproject_config(args: Namespace) -> Namespace: "disable_telemetry", "disable_imports_sorting", "git_remote", - "benchmarks_root" ] for key in supported_keys: if key in pyproject_config and ( diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index cd0bfc50a..6e4f744d7 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -363,7 +363,6 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None: for decorator in body_node.decorator_list ): self.is_staticmethod = True - print(f"static method found: {self.function_name}") return elif self.line_no: # If we have line number info, check if class has a static method with the same line number diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 52d1e4285..e046cf910 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -78,6 +78,15 @@ class BestOptimization(BaseModel): winning_benchmarking_test_results: TestResults winning_replay_benchmarking_test_results : Optional[TestResults] = None +@dataclass(frozen=True) +class BenchmarkKey: + file_name: str + function_name: str + line_number: int + + def __str__(self) -> str: + return f"{self.file_name}::{self.function_name}::{self.line_number}" + @dataclass class BenchmarkDetail: benchmark_name: str @@ -156,6 +165,7 @@ class OptimizedCandidateResult(BaseModel): best_test_runtime: int behavior_test_results: TestResults benchmarking_test_results: TestResults + replay_benchmarking_test_results: Optional[TestResults] = None optimization_candidate_index: int total_candidate_timing: int @@ -260,6 +270,7 @@ class FunctionParent: class OriginalCodeBaseline(BaseModel): behavioral_test_results: TestResults benchmarking_test_results: TestResults + replay_benchmarking_test_results: Optional[TestResults] = None runtime: int coverage_results: Optional[CoverageData] diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 7474f6991..fe82e5263 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -140,8 +140,6 @@ def optimize_function(self) -> Result[BestOptimization, str]: return Failure("Codeflash does not support async functions in the code to optimize.") code_print(code_context.read_writable_code) - logger.info("Read only code") - code_print(code_context.read_only_context_code) generated_test_paths = [ get_test_file_path( self.test_cfg.tests_root, self.function_to_optimize.function_name, test_index, test_type="unit" @@ -430,8 +428,8 @@ def determine_best_candidate( tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") if self.args.benchmark: - original_code_replay_runtime = original_code_baseline.benchmarking_test_results.total_replay_test_runtime() - candidate_replay_runtime = candidate_result.benchmarking_test_results.total_replay_test_runtime() + original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results.total_passed_runtime() + candidate_replay_runtime = candidate_result.replay_benchmarking_test_results.total_passed_runtime() replay_perf_gain = performance_gain( original_runtime_ns=original_code_replay_runtime, optimized_runtime_ns=candidate_replay_runtime, @@ -900,12 +898,14 @@ def establish_original_code_baseline( logger.debug(f"Total original code runtime (ns): {total_timing}") if self.args.benchmark: - logger.info(f"Total replay test runtime: {humanize_runtime(benchmarking_results.total_replay_test_runtime())}") + replay_benchmarking_test_results = benchmarking_results.filter(TestType.REPLAY_TEST) + logger.info(f"Total replay test runtime: {humanize_runtime(replay_benchmarking_test_results.total_passed_runtime())}") return Success( ( OriginalCodeBaseline( behavioral_test_results=behavioral_results, benchmarking_test_results=benchmarking_results, + replay_benchmarking_test_results = replay_benchmarking_test_results if self.args.benchmark else None, runtime=total_timing, coverage_results=coverage_results, ), @@ -1020,13 +1020,9 @@ def run_optimized_candidate( logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") if self.args.benchmark: - total_candidate_replay_timing = ( - candidate_benchmarking_results.total_replay_test_runtime() - if candidate_benchmarking_results - else 0 - ) + candidate_replay_benchmarking_results = candidate_benchmarking_results.filter(TestType.REPLAY_TEST) logger.debug( - f"Total optimized code {optimization_candidate_index} replay benchmark runtime (ns): {total_candidate_replay_timing}" + f"Total optimized code {optimization_candidate_index} replay benchmark runtime (ns): {candidate_replay_benchmarking_results.total_passed_runtime()}" ) return Success( OptimizedCandidateResult( @@ -1034,6 +1030,7 @@ def run_optimized_candidate( best_test_runtime=total_candidate_timing, behavior_test_results=candidate_behavior_results, benchmarking_test_results=candidate_benchmarking_results, + replay_benchmarking_test_results = candidate_replay_benchmarking_results if self.args.benchmark else None, optimization_candidate_index=optimization_candidate_index, total_candidate_timing=total_candidate_timing, ) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 8be62d963..4d17a5255 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient -from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils +from codeflash.benchmarking.plugin.plugin import CodeFlashBenchmarkPlugin from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table @@ -20,7 +20,7 @@ from codeflash.discovery.discover_unit_tests import discover_unit_tests from codeflash.discovery.functions_to_optimize import get_functions_to_optimize from codeflash.either import is_successful -from codeflash.models.models import ValidCode +from codeflash.models.models import ValidCode, BenchmarkKey from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.telemetry.posthog_cf import ph from codeflash.verification.test_results import TestType @@ -96,8 +96,8 @@ def run(self) -> None: project_root=self.args.project_root, module_root=self.args.module_root, ) - function_benchmark_timings = None - total_benchmark_timings = None + function_benchmark_timings: dict[str, dict[BenchmarkKey, int]] = {} + total_benchmark_timings: dict[BenchmarkKey, int] = {} if self.args.benchmark: with progress_bar( f"Running benchmarks in {self.args.benchmarks_root}", @@ -109,9 +109,7 @@ def run(self) -> None: 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: - instrument_codeflash_trace_decorator(fto) + instrument_codeflash_trace_decorator(file_to_funcs_to_optimize) trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" replay_tests_dir = Path(self.args.tests_root) / "codeflash_replay_tests" trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Run all tests that use pytest-benchmark @@ -119,8 +117,8 @@ def run(self) -> None: if replay_count == 0: logger.info(f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization") else: - function_benchmark_timings = BenchmarkDatabaseUtils.get_function_benchmark_timings(trace_file) - total_benchmark_timings = BenchmarkDatabaseUtils.get_benchmark_timings(trace_file) + function_benchmark_timings = CodeFlashBenchmarkPlugin.get_function_benchmark_timings(trace_file) + total_benchmark_timings = CodeFlashBenchmarkPlugin.get_benchmark_timings(trace_file) function_to_results = validate_and_format_benchmark_table(function_benchmark_timings, total_benchmark_timings) print_benchmark_table(function_to_results) logger.info("Finished tracing existing benchmarks") diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py index 7700c9cc1..6f4f397ab 100644 --- a/codeflash/verification/test_results.py +++ b/codeflash/verification/test_results.py @@ -1,6 +1,7 @@ from __future__ import annotations import sys +from collections import defaultdict from enum import Enum from pathlib import Path from typing import TYPE_CHECKING, Optional, cast @@ -124,6 +125,15 @@ def merge(self, other: TestResults) -> None: raise ValueError(msg) self.test_result_idx[k] = v + original_len + def filter(self, test_type: TestType) -> TestResults: + filtered_test_results = [] + filtered_test_results_idx = {} + for test_result in self.test_results: + if test_result.test_type == test_type: + filtered_test_results_idx[test_result.unique_invocation_loop_id] = len(filtered_test_results) + filtered_test_results.append(test_result) + return TestResults(test_results=filtered_test_results, test_result_idx=filtered_test_results_idx) + def get_by_unique_invocation_loop_id(self, unique_invocation_loop_id: str) -> FunctionTestInvocation | None: try: return self.test_results[self.test_result_idx[unique_invocation_loop_id]] @@ -174,22 +184,21 @@ def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree: return tree def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: + usable_runtime_by_id = defaultdict(list) for result in self.test_results: - if result.did_pass and not result.runtime: - msg = ( - f"Ignoring test case that passed but had no runtime -> {result.id}, " - f"Loop # {result.loop_index}, Test Type: {result.test_type}, " - f"Verification Type: {result.verification_type}" - ) - logger.debug(msg) - - usable_runtimes = [ - (result.id, result.runtime) for result in self.test_results if result.did_pass and result.runtime - ] - return { - usable_id: [runtime[1] for runtime in usable_runtimes if runtime[0] == usable_id] - for usable_id in {runtime[0] for runtime in usable_runtimes} - } + if result.did_pass: + if not result.runtime: + msg = ( + f"Ignoring test case that passed but had no runtime -> {result.id}, " + f"Loop # {result.loop_index}, Test Type: {result.test_type}, " + f"Verification Type: {result.verification_type}" + ) + logger.debug(msg) + else: + usable_runtime_by_id[result.id].append(result.runtime) + + return usable_runtime_by_id + def total_passed_runtime(self) -> int: """Calculate the sum of runtimes of all test cases that passed. @@ -202,35 +211,6 @@ def total_passed_runtime(self) -> int: [min(usable_runtime_data) for _, usable_runtime_data in self.usable_runtime_data_by_test_case().items()] ) - def usable_replay_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: - """Collect runtime data for replay tests that passed and have runtime information. - - :return: A dictionary mapping invocation IDs to lists of runtime values. - """ - usable_runtimes = [ - (result.id, result.runtime) - for result in self.test_results - if result.did_pass and result.runtime and result.test_type == TestType.REPLAY_TEST - ] - - return { - usable_id: [runtime[1] for runtime in usable_runtimes if runtime[0] == usable_id] - for usable_id in {runtime[0] for runtime in usable_runtimes} - } - - def total_replay_test_runtime(self) -> int: - """Calculate the sum of runtimes of replay test cases that passed, where a testcase runtime - is the minimum value of all looped execution runtimes. - - :return: The runtime in nanoseconds. - """ - replay_runtime_data = self.usable_replay_runtime_data_by_test_case() - - return sum([ - min(runtimes) - for invocation_id, runtimes in replay_runtime_data.items() - ]) if replay_runtime_data else 0 - def __iter__(self) -> Iterator[FunctionTestInvocation]: return iter(self.test_results) diff --git a/tests/test_instrument_codeflash_trace.py b/tests/test_instrument_codeflash_trace.py index 967d5d6f0..6b884c631 100644 --- a/tests/test_instrument_codeflash_trace.py +++ b/tests/test_instrument_codeflash_trace.py @@ -1,9 +1,10 @@ from __future__ import annotations +import tempfile from pathlib import Path -from codeflash.benchmarking.instrument_codeflash_trace import add_codeflash_decorator_to_code - +from codeflash.benchmarking.instrument_codeflash_trace import add_codeflash_decorator_to_code, \ + instrument_codeflash_trace_decorator from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize @@ -22,7 +23,7 @@ def normal_function(): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -34,6 +35,7 @@ def normal_function(): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_normal_method() -> None: """Test adding decorator to a normal method.""" code = """ @@ -50,7 +52,7 @@ def normal_method(self): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -63,6 +65,7 @@ def normal_method(self): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_classmethod() -> None: """Test adding decorator to a classmethod.""" code = """ @@ -80,7 +83,7 @@ def class_method(cls): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -94,6 +97,7 @@ def class_method(cls): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_staticmethod() -> None: """Test adding decorator to a staticmethod.""" code = """ @@ -111,7 +115,7 @@ def static_method(): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -125,6 +129,7 @@ def static_method(): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_init_function() -> None: """Test adding decorator to an __init__ function.""" code = """ @@ -141,7 +146,7 @@ def __init__(self, value): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -154,6 +159,7 @@ def __init__(self, value): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_with_multiple_decorators() -> None: """Test adding decorator to a function with multiple existing decorators.""" code = """ @@ -172,7 +178,7 @@ def property_method(self): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -187,6 +193,7 @@ def property_method(self): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_function_in_multiple_classes() -> None: """Test that only the right class's method gets the decorator.""" code = """ @@ -207,7 +214,7 @@ def test_method(self): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -224,6 +231,7 @@ def test_method(self): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_nonexistent_function() -> None: """Test that code remains unchanged when function doesn't exist.""" code = """ @@ -239,8 +247,223 @@ def existing_function(): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) # Code should remain unchanged assert modified_code.strip() == code.strip() + + +def test_add_decorator_to_multiple_functions() -> None: + """Test adding decorator to multiple functions.""" + code = """ +def function_one(): + return "First function" + +class TestClass: + def method_one(self): + return "First method" + + def method_two(self): + return "Second method" + +def function_two(): + return "Second function" +""" + + functions_to_optimize = [ + FunctionToOptimize( + function_name="function_one", + file_path=Path("dummy_path.py"), + parents=[] + ), + FunctionToOptimize( + function_name="method_two", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ), + FunctionToOptimize( + function_name="function_two", + file_path=Path("dummy_path.py"), + parents=[] + ) + ] + + modified_code = add_codeflash_decorator_to_code( + code=code, + functions_to_optimize=functions_to_optimize + ) + + expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +@codeflash_trace +def function_one(): + return "First function" + +class TestClass: + def method_one(self): + return "First method" + + @codeflash_trace + def method_two(self): + return "Second method" + +@codeflash_trace +def function_two(): + return "Second function" +""" + + assert modified_code.strip() == expected_code.strip() + + +def test_instrument_codeflash_trace_decorator_single_file() -> None: + """Test instrumenting codeflash trace decorator on a single file.""" + with tempfile.TemporaryDirectory() as temp_dir: + # Create a test Python file + test_file_path = Path(temp_dir) / "test_module.py" + test_file_content = """ +def function_one(): + return "First function" + +class TestClass: + def method_one(self): + return "First method" + + def method_two(self): + return "Second method" + +def function_two(): + return "Second function" +""" + test_file_path.write_text(test_file_content, encoding="utf-8") + + # Define functions to optimize + functions_to_optimize = [ + FunctionToOptimize( + function_name="function_one", + file_path=test_file_path, + parents=[] + ), + FunctionToOptimize( + function_name="method_two", + file_path=test_file_path, + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + ] + + # Execute the function being tested + instrument_codeflash_trace_decorator({test_file_path: functions_to_optimize}) + + # Read the modified file + modified_content = test_file_path.read_text(encoding="utf-8") + + # Define expected content (with isort applied) + expected_content = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace + + +@codeflash_trace +def function_one(): + return "First function" + +class TestClass: + def method_one(self): + return "First method" + + @codeflash_trace + def method_two(self): + return "Second method" + +def function_two(): + return "Second function" +""" + + # Compare the modified content with expected content + assert modified_content.strip() == expected_content.strip() + + +def test_instrument_codeflash_trace_decorator_multiple_files() -> None: + """Test instrumenting codeflash trace decorator on multiple files.""" + with tempfile.TemporaryDirectory() as temp_dir: + # Create first test Python file + test_file_1_path = Path(temp_dir) / "module_a.py" + test_file_1_content = """ +def function_a(): + return "Function in module A" + +class ClassA: + def method_a(self): + return "Method in ClassA" +""" + test_file_1_path.write_text(test_file_1_content, encoding="utf-8") + + # Create second test Python file + test_file_2_path = Path(temp_dir) / "module_b.py" + test_file_2_content =""" +def function_b(): + return "Function in module B" + +class ClassB: + @staticmethod + def static_method_b(): + return "Static method in ClassB" +""" + test_file_2_path.write_text(test_file_2_content, encoding="utf-8") + + # Define functions to optimize + file_to_funcs_to_optimize = { + test_file_1_path: [ + FunctionToOptimize( + function_name="function_a", + file_path=test_file_1_path, + parents=[] + ) + ], + test_file_2_path: [ + FunctionToOptimize( + function_name="static_method_b", + file_path=test_file_2_path, + parents=[FunctionParent(name="ClassB", type="ClassDef")] + ) + ] + } + + # Execute the function being tested + instrument_codeflash_trace_decorator(file_to_funcs_to_optimize) + + # Read the modified files + modified_content_1 = test_file_1_path.read_text(encoding="utf-8") + modified_content_2 = test_file_2_path.read_text(encoding="utf-8") + + # Define expected content for first file (with isort applied) + expected_content_1 = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace + + +@codeflash_trace +def function_a(): + return "Function in module A" + +class ClassA: + def method_a(self): + return "Method in ClassA" +""" + + # Define expected content for second file (with isort applied) + expected_content_2 = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace + + +def function_b(): + return "Function in module B" + +class ClassB: + @staticmethod + @codeflash_trace + def static_method_b(): + return "Static method in ClassB" +""" + + # Compare the modified content with expected content + assert modified_content_1.strip() == expected_content_1.strip() + assert modified_content_2.strip() == expected_content_2.strip() \ No newline at end of file diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index fcc5b0f67..f08b2485a 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,6 +1,6 @@ import sqlite3 -from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils +from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.replay_test import generate_replay_test from pathlib import Path @@ -27,7 +27,7 @@ def test_trace_benchmarks(): # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM function_calls ORDER BY benchmark_file_name, benchmark_function_name, function_name") + "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_name, benchmark_function_name, function_name") function_calls = cursor.fetchall() # Assert the length of function calls @@ -154,7 +154,6 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): finally: # cleanup shutil.rmtree(tests_root) - pass def test_trace_multithreaded_benchmark() -> None: project_root = Path(__file__).parent.parent / "code_to_optimize" @@ -173,13 +172,13 @@ def test_trace_multithreaded_benchmark() -> None: # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM function_calls ORDER BY benchmark_file_name, benchmark_function_name, function_name") + "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_name, benchmark_function_name, function_name") function_calls = cursor.fetchall() # Assert the length of function calls assert len(function_calls) == 10, f"Expected 10 function calls, but got {len(function_calls)}" - function_benchmark_timings = BenchmarkDatabaseUtils.get_function_benchmark_timings(output_file) - total_benchmark_timings = BenchmarkDatabaseUtils.get_benchmark_timings(output_file) + 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 @@ -208,5 +207,4 @@ def test_trace_multithreaded_benchmark() -> None: finally: # cleanup - shutil.rmtree(tests_root) - pass \ No newline at end of file + shutil.rmtree(tests_root) \ No newline at end of file From 57b80ec85f42f7b22b23e62969766d4c37bbe320 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 28 Mar 2025 14:58:49 -0700 Subject: [PATCH 031/122] started implementing group by benchmark --- codeflash/benchmarking/plugin/plugin.py | 11 ++---- codeflash/benchmarking/replay_test.py | 5 +-- codeflash/benchmarking/utils.py | 2 +- codeflash/models/models.py | 5 +-- codeflash/optimization/function_optimizer.py | 40 +++++++++++--------- codeflash/optimization/optimizer.py | 9 +++-- codeflash/verification/test_results.py | 20 +++++++++- 7 files changed, 56 insertions(+), 36 deletions(-) diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index 9d7da6ef2..a4805cca3 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -101,8 +101,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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func) # Initialize the inner dictionary if needed if qualified_name not in result: result[qualified_name] = {} @@ -152,8 +151,7 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]: # Process overhead information for row in cursor.fetchall(): benchmark_file, benchmark_func, benchmark_line, total_overhead_ns = row - benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + benchmark_key = BenchmarkKey(file_name=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 @@ -167,8 +165,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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + benchmark_key = BenchmarkKey(file_name=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 @@ -239,7 +236,7 @@ def test_something(benchmark): The return value of the function """ - benchmark_file_name = self.request.node.fspath.basename + benchmark_file_name = self.request.node.fspath benchmark_function_name = self.request.node.name line_number = int(str(sys._getframe(1).f_lineno)) # 1 frame up in the call stack diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 670d6e4bd..58bae35f1 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -196,12 +196,11 @@ def create_trace_replay_test_code( return imports + "\n" + metadata + "\n" + test_template def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework: str = "pytest", max_run_count: int = 100) -> int: - """Generate multiple replay tests from the traced function calls, grouping by benchmark name. + """Generate multiple replay tests from the traced function calls, grouped by benchmark. Args: trace_file_path: Path to the SQLite database file output_dir: Directory to write the generated tests (if None, only returns the code) - project_root: Root directory of the project for module imports test_framework: 'pytest' or 'unittest' max_run_count: Maximum number of runs to include per function @@ -267,7 +266,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework # Write to file if requested if output_dir: output_file = get_test_file_path( - test_dir=Path(output_dir), function_name=f"{benchmark_file_name[5:]}_{benchmark_function_name}", test_type="replay" + test_dir=Path(output_dir), function_name=f"{benchmark_file_name}_{benchmark_function_name}", test_type="replay" ) # Write test code to file, parents = true output_dir.mkdir(parents=True, exist_ok=True) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 5f14f141f..feb9ed0fc 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -93,7 +93,7 @@ def process_benchmark_data( for benchmark_key, og_benchmark_timing in fto_benchmark_timings.items(): try: - benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") + benchmark_file_name, benchmark_test_function = benchmark_key.split("::") except ValueError: continue # Skip malformed benchmark keys diff --git a/codeflash/models/models.py b/codeflash/models/models.py index e046cf910..f62131e2a 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -82,10 +82,9 @@ class BestOptimization(BaseModel): class BenchmarkKey: file_name: str function_name: str - line_number: int def __str__(self) -> str: - return f"{self.file_name}::{self.function_name}::{self.line_number}" + return f"{self.file_name}::{self.function_name}" @dataclass class BenchmarkDetail: @@ -270,7 +269,7 @@ class FunctionParent: class OriginalCodeBaseline(BaseModel): behavioral_test_results: TestResults benchmarking_test_results: TestResults - replay_benchmarking_test_results: Optional[TestResults] = None + replay_benchmarking_test_results: Optional[dict[BenchmarkKey, TestResults]] = None runtime: int coverage_results: Optional[CoverageData] diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index fe82e5263..435b5d0ad 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -88,8 +88,8 @@ def __init__( function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_ast: ast.FunctionDef | None = None, aiservice_client: AiServiceClient | None = None, - function_benchmark_timings: dict[str, int] | None = None, - total_benchmark_timings: dict[str, int] | None = None, + function_benchmark_timings: dict[BenchmarkKey, int] | None = None, + total_benchmark_timings: dict[BenchmarkKey, int] | None = None, args: Namespace | None = None, ) -> None: self.project_root = test_cfg.project_root_path @@ -428,20 +428,24 @@ def determine_best_candidate( tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") if self.args.benchmark: - original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results.total_passed_runtime() - candidate_replay_runtime = candidate_result.replay_benchmarking_test_results.total_passed_runtime() - replay_perf_gain = performance_gain( - original_runtime_ns=original_code_replay_runtime, - optimized_runtime_ns=candidate_replay_runtime, - ) - tree.add(f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") - tree.add( - f"Best benchmark replay runtime: {humanize_runtime(candidate_replay_runtime)} " - f"(measured over {candidate_result.max_loop_count} " - f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" - ) - tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain + 1:.1f}X") + + benchmark_keys = {(benchmark.file_name, benchmark.function_name) for benchmark in self.total_benchmark_timings} + test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmark(benchmark_keys) + for benchmark_key, test_results in test_results_by_benchmark.items(): + original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results[benchmark_key].total_passed_runtime() + candidate_replay_runtime = candidate_result.replay_benchmarking_test_results.total_passed_runtime() + replay_perf_gain = performance_gain( + original_runtime_ns=original_code_replay_runtime, + optimized_runtime_ns=candidate_replay_runtime, + ) + tree.add(f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") + tree.add( + f"Best benchmark replay runtime: {humanize_runtime(candidate_replay_runtime)} " + f"(measured over {candidate_result.max_loop_count} " + f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + ) + tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain + 1:.1f}X") best_optimization = BestOptimization( candidate=candidate, helper_functions=code_context.helper_functions, @@ -898,7 +902,7 @@ def establish_original_code_baseline( logger.debug(f"Total original code runtime (ns): {total_timing}") if self.args.benchmark: - replay_benchmarking_test_results = benchmarking_results.filter(TestType.REPLAY_TEST) + replay_benchmarking_test_results = benchmarking_results.filter_by_test_type(TestType.REPLAY_TEST) logger.info(f"Total replay test runtime: {humanize_runtime(replay_benchmarking_test_results.total_passed_runtime())}") return Success( ( @@ -1020,7 +1024,7 @@ def run_optimized_candidate( logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") if self.args.benchmark: - candidate_replay_benchmarking_results = candidate_benchmarking_results.filter(TestType.REPLAY_TEST) + candidate_replay_benchmarking_results = candidate_benchmarking_results.filter_by_test_type(TestType.REPLAY_TEST) logger.debug( f"Total optimized code {optimization_candidate_index} replay benchmark runtime (ns): {candidate_replay_benchmarking_results.total_passed_runtime()}" ) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 4d17a5255..35d91a274 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -60,8 +60,8 @@ def create_function_optimizer( function_to_optimize_ast: ast.FunctionDef | None = None, function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_source_code: str | None = "", - function_benchmark_timings: dict[str, dict[str, float]] | None = None, - total_benchmark_timings: dict[str, float] | None = None, + function_benchmark_timings: dict[str, dict[BenchmarkKey, float]] | None = None, + total_benchmark_timings: dict[BenchmarkKey, float] | None = None, ) -> FunctionOptimizer: return FunctionOptimizer( function_to_optimize=function_to_optimize, @@ -111,7 +111,10 @@ def run(self) -> None: try: instrument_codeflash_trace_decorator(file_to_funcs_to_optimize) trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" - replay_tests_dir = Path(self.args.tests_root) / "codeflash_replay_tests" + if trace_file.exists(): + trace_file.unlink() + + replay_tests_dir = Path(self.args.tests_root) trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Run all tests that use pytest-benchmark replay_count = generate_replay_test(trace_file, replay_tests_dir) if replay_count == 0: diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py index 6f4f397ab..25f258e26 100644 --- a/codeflash/verification/test_results.py +++ b/codeflash/verification/test_results.py @@ -125,7 +125,7 @@ def merge(self, other: TestResults) -> None: raise ValueError(msg) self.test_result_idx[k] = v + original_len - def filter(self, test_type: TestType) -> TestResults: + def filter_by_test_type(self, test_type: TestType) -> TestResults: filtered_test_results = [] filtered_test_results_idx = {} for test_result in self.test_results: @@ -134,6 +134,24 @@ def filter(self, test_type: TestType) -> TestResults: filtered_test_results.append(test_result) return TestResults(test_results=filtered_test_results, test_result_idx=filtered_test_results_idx) + def group_by_benchmark(self, benchmark_key_set:set[tuple[str,str]]) -> dict[tuple[str,str],TestResults]: + """Group TestResults by benchmark key. + + For now, use a tuple of (file_path, function_name) as the benchmark key. Can't import BenchmarkKey because of circular import. + + Args: + benchmark_key_set (set[tuple[str,str]]): A set of tuples of (file_path, function_name) + + Returns: + TestResults: A new TestResults object with the test results grouped by benchmark key. + + """ + test_result_by_benchmark = defaultdict(TestResults) + for test_result in self.test_results: + if test_result.test_type == TestType.REPLAY_TEST and (test_result.id.test_module_path,test_result.id.test_function_name) in benchmark_key_set: + test_result_by_benchmark[(test_result.id.test_module_path,test_result.id.test_function_name)].add(test_result) + return test_result_by_benchmark + def get_by_unique_invocation_loop_id(self, unique_invocation_loop_id: str) -> FunctionTestInvocation | None: try: return self.test_results[self.test_result_idx[unique_invocation_loop_id]] From 56e34474a90837dce42c9c72d1a403d7d4eb4371 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Mon, 31 Mar 2025 16:47:58 -0700 Subject: [PATCH 032/122] reworked matching benchmark key to test results. --- codeflash/benchmarking/codeflash_trace.py | 12 ++-- codeflash/benchmarking/plugin/plugin.py | 24 ++++---- codeflash/benchmarking/replay_test.py | 64 +++++++++++--------- codeflash/benchmarking/utils.py | 32 ++++------ codeflash/github/PrComment.py | 1 - codeflash/models/models.py | 61 +++++++++++-------- codeflash/optimization/function_optimizer.py | 35 +++++------ codeflash/optimization/optimizer.py | 7 ++- codeflash/result/explanation.py | 1 - tests/test_trace_benchmarks.py | 42 ++++++------- tests/test_unit_test_discovery.py | 1 - 11 files changed, 144 insertions(+), 136 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 2ae57307b..8c307f8a0 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -33,8 +33,8 @@ def setup(self, trace_path: str) -> None: cur.execute("PRAGMA synchronous = OFF") cur.execute( "CREATE TABLE IF NOT EXISTS benchmark_function_timings(" - "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," - "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," + "function_name TEXT, class_name TEXT, module_name TEXT, file_path TEXT," + "benchmark_function_name TEXT, benchmark_file_path TEXT, benchmark_line_number INTEGER," "function_time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" ) self._connection.commit() @@ -62,8 +62,8 @@ def write_function_timings(self) -> None: # Insert data into the benchmark_function_timings table cur.executemany( "INSERT INTO benchmark_function_timings" - "(function_name, class_name, module_name, file_name, benchmark_function_name, " - "benchmark_file_name, benchmark_line_number, function_time_ns, overhead_time_ns, args, kwargs) " + "(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) " "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", self.function_calls_data ) @@ -115,7 +115,7 @@ def wrapper(*args, **kwargs): # 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", "") + benchmark_file_path = os.environ.get("CODEFLASH_BENCHMARK_FILE_PATH", "") benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "") # Get class name class_name = "" @@ -151,7 +151,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_name, benchmark_line_number, execution_time, + benchmark_function_name, benchmark_file_path, benchmark_line_number, execution_time, overhead_time, pickled_args, pickled_kwargs) ) return result diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index a4805cca3..09858601c 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -24,7 +24,7 @@ def setup(self, trace_path:str) -> None: cur.execute("PRAGMA synchronous = OFF") cur.execute( "CREATE TABLE IF NOT EXISTS benchmark_timings(" - "benchmark_file_name TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER," + "benchmark_file_path TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER," "benchmark_time_ns INTEGER)" ) self._connection.commit() @@ -47,7 +47,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_name, benchmark_function_name, benchmark_line_number, benchmark_time_ns) VALUES (?, ?, ?, ?)", + "INSERT INTO benchmark_timings (benchmark_file_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns) VALUES (?, ?, ?, ?)", self.benchmark_timings ) self._connection.commit() @@ -86,7 +86,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_name, benchmark_function_name, benchmark_line_number, function_time_ns " + "benchmark_file_path, benchmark_function_name, benchmark_line_number, function_time_ns " "FROM benchmark_function_timings" ) @@ -101,7 +101,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_name=benchmark_file, function_name=benchmark_func) + benchmark_key = BenchmarkKey(file_path=benchmark_file, function_name=benchmark_func) # Initialize the inner dictionary if needed if qualified_name not in result: result[qualified_name] = {} @@ -143,20 +143,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_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " + "SELECT benchmark_file_path, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " "FROM benchmark_function_timings " - "GROUP BY benchmark_file_name, benchmark_function_name, benchmark_line_number" + "GROUP BY benchmark_file_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_name=benchmark_file, function_name=benchmark_func) + benchmark_key = BenchmarkKey(file_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_name, benchmark_function_name, benchmark_line_number, benchmark_time_ns " + "SELECT benchmark_file_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns " "FROM benchmark_timings" ) @@ -165,7 +165,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_name=benchmark_file, function_name=benchmark_func) + benchmark_key = BenchmarkKey(file_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 @@ -236,13 +236,13 @@ def test_something(benchmark): The return value of the function """ - benchmark_file_name = self.request.node.fspath + benchmark_file_path = str(self.request.node.fspath) 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_NAME"] = benchmark_file_name + os.environ["CODEFLASH_BENCHMARK_FILE_PATH"] = benchmark_file_path os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(line_number) os.environ["CODEFLASH_BENCHMARKING"] = "True" @@ -260,7 +260,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_name, benchmark_function_name, line_number, end - start)) + (benchmark_file_path, benchmark_function_name, line_number, end - start)) return result diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 58bae35f1..9ecac2ec4 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -2,18 +2,21 @@ import sqlite3 import textwrap -from collections.abc import Generator -from typing import Any, Dict +from pathlib import Path +from typing import TYPE_CHECKING, Any import isort from codeflash.cli_cmds.console import logger from codeflash.discovery.functions_to_optimize import inspect_top_level_functions_or_methods from codeflash.verification.verification_utils import get_test_file_path -from pathlib import Path + +if TYPE_CHECKING: + from collections.abc import Generator + def get_next_arg_and_return( - trace_file: str, function_name: str, file_name: str, class_name: str | None = None, num_to_get: int = 256 + trace_file: str, function_name: str, file_path: str, class_name: str | None = None, num_to_get: int = 256 ) -> Generator[Any]: db = sqlite3.connect(trace_file) cur = db.cursor() @@ -21,13 +24,13 @@ def get_next_arg_and_return( if class_name is not None: cursor = cur.execute( - "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_name = ? AND class_name = ? LIMIT ?", - (function_name, file_name, class_name, limit), + "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_path = ? AND class_name = ? LIMIT ?", + (function_name, file_path, class_name, limit), ) else: cursor = cur.execute( - "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_name = ? AND class_name = '' LIMIT ?", - (function_name, file_name, limit), + "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_path = ? AND class_name = '' LIMIT ?", + (function_name, file_path, limit), ) while (val := cursor.fetchone()) is not None: @@ -88,7 +91,7 @@ def create_trace_replay_test_code( # Templates for different types of tests test_function_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = {function_name}(*args, **kwargs) @@ -97,7 +100,7 @@ def create_trace_replay_test_code( test_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} function_name = "{orig_function_name}" @@ -112,7 +115,7 @@ def create_trace_replay_test_code( test_class_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} if not args: @@ -122,7 +125,7 @@ def create_trace_replay_test_code( ) test_static_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} ret = {class_name_alias}{method_name}(*args, **kwargs) @@ -140,13 +143,13 @@ def create_trace_replay_test_code( module_name = func.get("module_name") function_name = func.get("function_name") class_name = func.get("class_name") - file_name = func.get("file_name") + file_path = func.get("file_path") function_properties = func.get("function_properties") if not class_name: alias = get_function_alias(module_name, function_name) test_body = test_function_body.format( function_name=alias, - file_name=file_name, + file_path=file_path, orig_function_name=function_name, max_run_count=max_run_count, ) @@ -160,7 +163,7 @@ def create_trace_replay_test_code( if function_properties.is_classmethod: test_body = test_class_method_body.format( orig_function_name=function_name, - file_name=file_name, + file_path=file_path, class_name_alias=class_name_alias, class_name=class_name, method_name=method_name, @@ -170,7 +173,7 @@ def create_trace_replay_test_code( elif function_properties.is_staticmethod: test_body = test_static_method_body.format( orig_function_name=function_name, - file_name=file_name, + file_path=file_path, class_name_alias=class_name_alias, class_name=class_name, method_name=method_name, @@ -180,7 +183,7 @@ def create_trace_replay_test_code( else: test_body = test_method_body.format( orig_function_name=function_name, - file_name=file_name, + file_path=file_path, class_name_alias=class_name_alias, class_name=class_name, method_name=method_name, @@ -216,42 +219,41 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework # Get distinct benchmark names cursor.execute( - "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM benchmark_function_timings" + "SELECT DISTINCT benchmark_function_name, benchmark_file_path FROM benchmark_function_timings" ) benchmarks = cursor.fetchall() # Generate a test for each benchmark for benchmark in benchmarks: - benchmark_function_name, benchmark_file_name = benchmark + benchmark_function_name, benchmark_file_path = benchmark # Get functions associated with this benchmark cursor.execute( - "SELECT DISTINCT function_name, class_name, module_name, file_name, benchmark_line_number FROM benchmark_function_timings " - "WHERE benchmark_function_name = ? AND benchmark_file_name = ?", - (benchmark_function_name, benchmark_file_name) + "SELECT DISTINCT function_name, class_name, module_name, file_path, benchmark_line_number FROM benchmark_function_timings " + "WHERE benchmark_function_name = ? AND benchmark_file_path = ?", + (benchmark_function_name, benchmark_file_path) ) functions_data = [] for func_row in cursor.fetchall(): - function_name, class_name, module_name, file_name, benchmark_line_number = func_row - + function_name, class_name, module_name, file_path, benchmark_line_number = func_row # Add this function to our list functions_data.append({ "function_name": function_name, "class_name": class_name, - "file_name": file_name, + "file_path": file_path, "module_name": module_name, "benchmark_function_name": benchmark_function_name, - "benchmark_file_name": benchmark_file_name, + "benchmark_file_path": benchmark_file_path, "benchmark_line_number": benchmark_line_number, "function_properties": inspect_top_level_functions_or_methods( - file_name=file_name, + file_name=Path(file_path), function_or_method_name=function_name, class_name=class_name, ) }) if not functions_data: - logger.info(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") + logger.info(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_path}") continue # Generate the test code for this benchmark @@ -265,17 +267,19 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework # Write to file if requested if output_dir: + name = Path(benchmark_file_path).name.split(".")[0][5:] # remove "test_" from the name since we add it in later output_file = get_test_file_path( - test_dir=Path(output_dir), function_name=f"{benchmark_file_name}_{benchmark_function_name}", test_type="replay" + test_dir=Path(output_dir), function_name=f"{name}_{benchmark_function_name}", test_type="replay" ) # Write test code to file, parents = true output_dir.mkdir(parents=True, exist_ok=True) output_file.write_text(test_code, "utf-8") count += 1 - logger.info(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") + logger.info(f"Replay test for benchmark `{benchmark_function_name}` in {name} written to {output_file}") conn.close() except Exception as e: logger.info(f"Error generating replay tests: {e}") + return count \ No newline at end of file diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index feb9ed0fc..1d8b22f50 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -11,7 +11,7 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[BenchmarkKey, int]], - total_benchmark_timings: dict[BenchmarkKey, int]) -> dict[str, list[tuple[str, float, float, float]]]: + total_benchmark_timings: dict[BenchmarkKey, int]) -> dict[str, list[tuple[BenchmarkKey, float, float, float]]]: function_to_result = {} # Process each function's benchmark data for func_path, test_times in function_benchmark_timings.items(): @@ -23,18 +23,18 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, di logger.debug(f"Skipping test {benchmark_key} due to func_time {func_time} > total_time {total_time}") # If the function time is greater than total time, likely to have multithreading / multiprocessing issues. # Do not try to project the optimization impact for this function. - sorted_tests.append((str(benchmark_key), 0.0, 0.0, 0.0)) + sorted_tests.append((benchmark_key, 0.0, 0.0, 0.0)) if total_time > 0: percentage = (func_time / total_time) * 100 # Convert nanoseconds to milliseconds func_time_ms = func_time / 1_000_000 total_time_ms = total_time / 1_000_000 - sorted_tests.append((str(benchmark_key), total_time_ms, func_time_ms, percentage)) + sorted_tests.append((benchmark_key, total_time_ms, func_time_ms, percentage)) sorted_tests.sort(key=lambda x: x[3], reverse=True) function_to_result[func_path] = sorted_tests return function_to_result -def print_benchmark_table(function_to_results: dict[str, list[tuple[str, float, float, float]]]) -> None: +def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey, float, float, float]]]) -> None: console = Console() for func_path, sorted_tests in function_to_results.items(): function_name = func_path.split(":")[-1] @@ -48,19 +48,17 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[str, float, table.add_column("Function Time (ms)", justify="right", style="yellow") table.add_column("Percentage (%)", justify="right", style="red") - for test_name, total_time, func_time, percentage in sorted_tests: - benchmark_file, benchmark_func, benchmark_line = test_name.split("::") - benchmark_name = f"{benchmark_file}::{benchmark_func}" + for benchmark_key, total_time, func_time, percentage in sorted_tests: if total_time == 0.0: table.add_row( - benchmark_name, + f"{benchmark_key.file_path}::{benchmark_key.function_name}", "N/A", "N/A", "N/A" ) else: table.add_row( - benchmark_name, + f"{benchmark_key.file_path}::{benchmark_key.function_name}", f"{total_time:.3f}", f"{func_time:.3f}", f"{percentage:.2f}" @@ -71,9 +69,9 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[str, float, def process_benchmark_data( - replay_performance_gain: float, - fto_benchmark_timings: dict[str, int], - total_benchmark_timings: dict[str, int] + replay_performance_gain: dict[BenchmarkKey, float], + fto_benchmark_timings: dict[BenchmarkKey, int], + total_benchmark_timings: dict[BenchmarkKey, int] ) -> Optional[ProcessedBenchmarkInfo]: """Process benchmark data and generate detailed benchmark information. @@ -92,10 +90,6 @@ def process_benchmark_data( benchmark_details = [] for benchmark_key, og_benchmark_timing in fto_benchmark_timings.items(): - try: - benchmark_file_name, benchmark_test_function = benchmark_key.split("::") - except ValueError: - continue # Skip malformed benchmark keys total_benchmark_timing = total_benchmark_timings.get(benchmark_key, 0) @@ -104,7 +98,7 @@ def process_benchmark_data( # Calculate expected new benchmark timing expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + ( - 1 / (replay_performance_gain + 1) + 1 / (replay_performance_gain[benchmark_key] + 1) ) * og_benchmark_timing # Calculate speedup @@ -112,8 +106,8 @@ def process_benchmark_data( benchmark_details.append( BenchmarkDetail( - benchmark_name=benchmark_file_name, - test_function=benchmark_test_function, + benchmark_name=benchmark_key.file_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)), speedup_percent=benchmark_speedup_percent diff --git a/codeflash/github/PrComment.py b/codeflash/github/PrComment.py index 8da239d38..1e66c5608 100644 --- a/codeflash/github/PrComment.py +++ b/codeflash/github/PrComment.py @@ -6,7 +6,6 @@ from codeflash.code_utils.time_utils import humanize_runtime from codeflash.models.models import BenchmarkDetail -from codeflash.verification.test_results import TestResults from codeflash.models.models import TestResults diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 6b69cbc52..ed0360eef 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections import defaultdict from typing import TYPE_CHECKING from rich.tree import Tree @@ -11,7 +12,7 @@ import enum import re import sys -from collections.abc import Collection, Iterator +from collections.abc import Collection from enum import Enum, IntEnum from pathlib import Path from re import Pattern @@ -22,10 +23,8 @@ from pydantic.dataclasses import dataclass from codeflash.cli_cmds.console import console, logger -from codeflash.code_utils.code_utils import validate_python_code +from codeflash.code_utils.code_utils import module_name_from_file_path, validate_python_code from codeflash.code_utils.env_utils import is_end_to_end -from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.verification.test_results import TestResults, TestType from codeflash.verification.comparator import comparator # If the method spam is in the class Ham, which is at the top level of the module eggs in the package foo, the fully @@ -74,19 +73,18 @@ class BestOptimization(BaseModel): candidate: OptimizedCandidate helper_functions: list[FunctionSource] runtime: int - replay_runtime: Optional[int] = None - replay_performance_gain: Optional[float] = None + replay_performance_gain: Optional[dict[BenchmarkKey,float]] = None winning_behavioral_test_results: TestResults winning_benchmarking_test_results: TestResults winning_replay_benchmarking_test_results : Optional[TestResults] = None @dataclass(frozen=True) class BenchmarkKey: - file_name: str + file_path: str function_name: str def __str__(self) -> str: - return f"{self.file_name}::{self.function_name}" + return f"{self.file_path}::{self.function_name}" @dataclass class BenchmarkDetail: @@ -166,7 +164,7 @@ class OptimizedCandidateResult(BaseModel): best_test_runtime: int behavior_test_results: TestResults benchmarking_test_results: TestResults - replay_benchmarking_test_results: Optional[TestResults] = None + replay_benchmarking_test_results: Optional[dict[BenchmarkKey, TestResults]] = None optimization_candidate_index: int total_candidate_timing: int @@ -473,6 +471,21 @@ def merge(self, other: TestResults) -> None: raise ValueError(msg) self.test_result_idx[k] = v + original_len + def group_by_benchmarks(self, benchmark_keys:list[BenchmarkKey], benchmark_replay_test_dir: Path, project_root: Path) -> dict[BenchmarkKey, TestResults]: + """Group TestResults by benchmark for calculating improvements for each benchmark.""" + + 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:]}_{benchmark_key.function_name}__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(): + if test_result.id.test_module_path.startswith(module_path): + test_results_by_benchmark[benchmark_key].add(test_result) + + return test_results_by_benchmark + def get_by_unique_invocation_loop_id(self, unique_invocation_loop_id: str) -> FunctionTestInvocation | None: try: return self.test_results[self.test_result_idx[unique_invocation_loop_id]] @@ -520,25 +533,23 @@ def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree: tree.add( f"{test_type.to_name()} - Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']}" ) - return tree + return def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: + + usable_runtime = defaultdict(list) for result in self.test_results: - if result.did_pass and not result.runtime: - msg = ( - f"Ignoring test case that passed but had no runtime -> {result.id}, " - f"Loop # {result.loop_index}, Test Type: {result.test_type}, " - f"Verification Type: {result.verification_type}" - ) - logger.debug(msg) - - usable_runtimes = [ - (result.id, result.runtime) for result in self.test_results if result.did_pass and result.runtime - ] - return { - usable_id: [runtime[1] for runtime in usable_runtimes if runtime[0] == usable_id] - for usable_id in {runtime[0] for runtime in usable_runtimes} - } + if result.did_pass: + if not result.runtime: + msg = ( + f"Ignoring test case that passed but had no runtime -> {result.id}, " + f"Loop # {result.loop_index}, Test Type: {result.test_type}, " + f"Verification Type: {result.verification_type}" + ) + logger.debug(msg) + else: + usable_runtime[result.id].append(result.runtime) + return usable_runtime def total_passed_runtime(self) -> int: """Calculate the sum of runtimes of all test cases that passed. diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 4111766f9..807fd3a8c 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -58,7 +58,7 @@ TestFiles, TestingMode, TestResults, - TestType, + TestType, BenchmarkKey, ) from codeflash.result.create_pr import check_create_pr, existing_tests_source_for from codeflash.result.critic import coverage_critic, performance_gain, quantity_of_tests_critic, speedup_critic @@ -428,30 +428,30 @@ def determine_best_candidate( ) tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") + replay_perf_gain = {} if self.args.benchmark: - - benchmark_keys = {(benchmark.file_name, benchmark.function_name) for benchmark in self.total_benchmark_timings} - test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmark(benchmark_keys) - for benchmark_key, test_results in test_results_by_benchmark.items(): + logger.info(f"Calculating benchmark improvement..") + test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) + for benchmark_key, candidate_test_results in test_results_by_benchmark.items(): original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results[benchmark_key].total_passed_runtime() - candidate_replay_runtime = candidate_result.replay_benchmarking_test_results.total_passed_runtime() - replay_perf_gain = performance_gain( + candidate_replay_runtime = candidate_test_results.total_passed_runtime() + replay_perf_gain[benchmark_key] = performance_gain( original_runtime_ns=original_code_replay_runtime, optimized_runtime_ns=candidate_replay_runtime, ) - tree.add(f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") + tree.add( + f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") tree.add( f"Best benchmark replay runtime: {humanize_runtime(candidate_replay_runtime)} " f"(measured over {candidate_result.max_loop_count} " f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" ) - tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain + 1:.1f}X") + tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain[benchmark_key] * 100:.1f}%") + tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain[benchmark_key] + 1:.1f}X") best_optimization = BestOptimization( candidate=candidate, helper_functions=code_context.helper_functions, runtime=best_test_runtime, - replay_runtime=candidate_replay_runtime if self.args.benchmark else None, winning_behavioral_test_results=candidate_result.behavior_test_results, replay_performance_gain=replay_perf_gain if self.args.benchmark else None, winning_benchmarking_test_results=candidate_result.benchmarking_test_results, @@ -903,8 +903,10 @@ def establish_original_code_baseline( logger.debug(f"Total original code runtime (ns): {total_timing}") if self.args.benchmark: - replay_benchmarking_test_results = benchmarking_results.filter_by_test_type(TestType.REPLAY_TEST) - logger.info(f"Total replay test runtime: {humanize_runtime(replay_benchmarking_test_results.total_passed_runtime())}") + replay_benchmarking_test_results = benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) + for benchmark_name, benchmark_results in replay_benchmarking_test_results.items(): + + logger.info(f"Replay benchmark '{benchmark_name}' runtime: {humanize_runtime(benchmark_results.total_passed_runtime())}") return Success( ( OriginalCodeBaseline( @@ -1025,10 +1027,9 @@ def run_optimized_candidate( logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") if self.args.benchmark: - candidate_replay_benchmarking_results = candidate_benchmarking_results.filter_by_test_type(TestType.REPLAY_TEST) - logger.debug( - f"Total optimized code {optimization_candidate_index} replay benchmark runtime (ns): {candidate_replay_benchmarking_results.total_passed_runtime()}" - ) + candidate_replay_benchmarking_results = candidate_benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) + for benchmark_name, benchmark_results in candidate_replay_benchmarking_results.items(): + logger.debug(f"Benchmark {benchmark_name} runtime (ns): {humanize_runtime(benchmark_results.total_passed_runtime())}") return Success( OptimizedCandidateResult( max_loop_count=loop_count, diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 86a0e6972..c0f2bf0ab 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -48,6 +48,7 @@ def __init__(self, args: Namespace) -> None: project_root_path=args.project_root, test_framework=args.test_framework, pytest_cmd=args.pytest_cmd, + benchmark_tests_root=args.benchmarks_root if "benchmark" in args and "benchmarks_root" in args else None, ) self.aiservice_client = AiServiceClient() @@ -114,7 +115,7 @@ def run(self) -> None: if trace_file.exists(): trace_file.unlink() - replay_tests_dir = Path(self.args.tests_root) + replay_tests_dir = Path(self.args.benchmarks_root) / "codeflash_replay_tests" trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Run all tests that use pytest-benchmark replay_count = generate_replay_test(trace_file, replay_tests_dir) if replay_count == 0: @@ -251,8 +252,8 @@ def run(self) -> None: if function_optimizer.test_cfg.concolic_test_root_dir: shutil.rmtree(function_optimizer.test_cfg.concolic_test_root_dir, ignore_errors=True) if self.args.benchmark: - if replay_tests_dir.exists(): - shutil.rmtree(replay_tests_dir, ignore_errors=True) + # if replay_tests_dir.exists(): + # shutil.rmtree(replay_tests_dir, ignore_errors=True) trace_file.unlink(missing_ok=True) if hasattr(get_run_tmp_file, "tmpdir"): get_run_tmp_file.tmpdir.cleanup() diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 7c8cf87e0..e56558a94 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -6,7 +6,6 @@ from codeflash.code_utils.time_utils import humanize_runtime from codeflash.models.models import BenchmarkDetail, TestResults -from codeflash.verification.test_results import TestResults @dataclass(frozen=True, config={"arbitrary_types_allowed": True}) diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index f08b2485a..c67c7e87d 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -27,7 +27,7 @@ def test_trace_benchmarks(): # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_name, benchmark_function_name, function_name") + "SELECT function_name, class_name, module_name, file_path, benchmark_function_name, benchmark_file_path, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_path, benchmark_function_name, function_name") function_calls = cursor.fetchall() # Assert the length of function calls @@ -39,44 +39,44 @@ def test_trace_benchmarks(): expected_calls = [ ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", "test_benchmark_bubble_sort.py", 20), + "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 20), ("sort_class", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", "test_benchmark_bubble_sort.py", 18), + "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 18), ("sort_static", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", "test_benchmark_bubble_sort.py", 19), + "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 19), ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", "test_benchmark_bubble_sort.py", 17), + "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 17), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_sort", "test_benchmark_bubble_sort.py", 7), + "test_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 7), ("compute_and_sort", "", "code_to_optimize.process_and_bubble_sort_codeflash_trace", f"{process_and_bubble_sort_path}", - "test_compute_and_sort", "test_process_and_sort.py", 4), + "test_compute_and_sort", str(benchmarks_root / "test_process_and_sort.py"), 4), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_no_func", "test_process_and_sort.py", 8), + "test_no_func", str(benchmarks_root / "test_process_and_sort.py"), 8), ] 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_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_file_name" + assert actual[5] == expected[5], f"Mismatch at index {idx} for benchmark_file_path" assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number" # Close connection conn.close() generate_replay_test(output_file, tests_root) - test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort_py_test_class_sort__replay_test_0.py") + test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort_test_class_sort__replay_test_0.py") assert test_class_sort_path.exists() test_class_sort_code = f""" import dill as pickle @@ -89,7 +89,7 @@ def test_trace_benchmarks(): trace_file_path = r"{output_file.as_posix()}" def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "sorter" @@ -102,7 +102,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): ret = instance.sorter(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) if not args: @@ -110,13 +110,13 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_class(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_static(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_static(*args, **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "__init__" @@ -131,7 +131,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): """ assert test_class_sort_path.read_text("utf-8").strip()==test_class_sort_code.strip() - test_sort_path = tests_root / Path("test_benchmark_bubble_sort_py_test_sort__replay_test_0.py") + test_sort_path = tests_root / Path("test_benchmark_bubble_sort_test_sort__replay_test_0.py") assert test_sort_path.exists() test_sort_code = f""" import dill as pickle @@ -144,7 +144,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): trace_file_path = r"{output_file}" def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"{bubble_sort_path}", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_path=r"{bubble_sort_path}", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_sorter(*args, **kwargs) @@ -172,7 +172,7 @@ def test_trace_multithreaded_benchmark() -> None: # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_name, benchmark_function_name, function_name") + "SELECT function_name, class_name, module_name, file_path, benchmark_function_name, benchmark_file_path, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_path, benchmark_function_name, function_name") function_calls = cursor.fetchall() # Assert the length of function calls @@ -192,15 +192,15 @@ def test_trace_multithreaded_benchmark() -> None: expected_calls = [ ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_benchmark_sort", "test_multithread_sort.py", 4), + "test_benchmark_sort", str(benchmarks_root / "test_multithread_sort.py"), 4), ] 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_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_file_name" + assert actual[5] == expected[5], f"Mismatch at index {idx} for benchmark_file_path" assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number" # Close connection conn.close() diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index c05b79e63..8c3bc35c8 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -3,7 +3,6 @@ from pathlib import Path from codeflash.discovery.discover_unit_tests import discover_unit_tests -from codeflash.verification.test_results import TestType from codeflash.verification.verification_utils import TestConfig From 5f86bdddb5e1b0c3060ab93571dc6783c434b0d3 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Mon, 31 Mar 2025 20:20:57 -0700 Subject: [PATCH 033/122] PRAGMA journal to memory to make it faster --- codeflash/benchmarking/codeflash_trace.py | 1 + codeflash/benchmarking/plugin/plugin.py | 1 + 2 files changed, 2 insertions(+) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 8c307f8a0..bcbb3268c 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -31,6 +31,7 @@ def setup(self, trace_path: str) -> None: self._connection = sqlite3.connect(self._trace_path) cur = self._connection.cursor() cur.execute("PRAGMA synchronous = OFF") + cur.execute("PRAGMA journal_mode = MEMORY") cur.execute( "CREATE TABLE IF NOT EXISTS benchmark_function_timings(" "function_name TEXT, class_name TEXT, module_name TEXT, file_path TEXT," diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index 09858601c..b022f9afb 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -22,6 +22,7 @@ def setup(self, trace_path:str) -> None: self._connection = sqlite3.connect(self._trace_path) cur = self._connection.cursor() cur.execute("PRAGMA synchronous = OFF") + 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," From 9764c25941010cc15633e425fb8b993f56150ac1 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 1 Apr 2025 10:53:14 -0700 Subject: [PATCH 034/122] benchmarks root must be subdir of tests root --- .../tests/unittest/test_bubble_sort.py | 36 +++++++++---------- .../unittest/test_bubble_sort_parametrized.py | 36 +++++++++---------- codeflash/cli_cmds/cli.py | 3 ++ 3 files changed, 39 insertions(+), 36 deletions(-) diff --git a/code_to_optimize/tests/unittest/test_bubble_sort.py b/code_to_optimize/tests/unittest/test_bubble_sort.py index 4c76414ef..200f82b7a 100644 --- a/code_to_optimize/tests/unittest/test_bubble_sort.py +++ b/code_to_optimize/tests/unittest/test_bubble_sort.py @@ -1,18 +1,18 @@ -# import unittest -# -# from code_to_optimize.bubble_sort import sorter -# -# -# class TestPigLatin(unittest.TestCase): -# def test_sort(self): -# input = [5, 4, 3, 2, 1, 0] -# output = sorter(input) -# self.assertEqual(output, [0, 1, 2, 3, 4, 5]) -# -# input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0] -# output = sorter(input) -# self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) -# -# input = list(reversed(range(5000))) -# output = sorter(input) -# self.assertEqual(output, list(range(5000))) +import unittest + +from code_to_optimize.bubble_sort import sorter + + +class TestPigLatin(unittest.TestCase): + def test_sort(self): + input = [5, 4, 3, 2, 1, 0] + output = sorter(input) + self.assertEqual(output, [0, 1, 2, 3, 4, 5]) + + input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0] + output = sorter(input) + self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) + + input = list(reversed(range(5000))) + output = sorter(input) + self.assertEqual(output, list(range(5000))) diff --git a/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py b/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py index c1aef993b..59c86abc8 100644 --- a/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py +++ b/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py @@ -1,18 +1,18 @@ -# import unittest -# -# from parameterized import parameterized -# -# from code_to_optimize.bubble_sort import sorter -# -# -# class TestPigLatin(unittest.TestCase): -# @parameterized.expand( -# [ -# ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]), -# ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), -# (list(reversed(range(50))), list(range(50))), -# ] -# ) -# def test_sort(self, input, expected_output): -# output = sorter(input) -# self.assertEqual(output, expected_output) +import unittest + +from parameterized import parameterized + +from code_to_optimize.bubble_sort import sorter + + +class TestPigLatin(unittest.TestCase): + @parameterized.expand( + [ + ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]), + ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), + (list(reversed(range(50))), list(range(50))), + ] + ) + def test_sort(self, input, expected_output): + output = sorter(input) + self.assertEqual(output, expected_output) diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 9dae8eea8..07652f707 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -135,6 +135,9 @@ def process_pyproject_config(args: Namespace) -> Namespace: if args.benchmark: assert args.benchmarks_root is not None, "--benchmarks-root must be specified when running with --benchmark" assert Path(args.benchmarks_root).is_dir(), f"--benchmarks-root {args.benchmarks_root} must be a valid directory" + assert Path(args.benchmarks_root).is_relative_to(Path(args.tests_root)), ( + f"--benchmarks-root {args.benchmarks_root} must be a subdirectory of --tests-root {args.tests_root}" + ) if env_utils.get_pr_number() is not None: assert env_utils.ensure_codeflash_api_key(), ( "Codeflash API key not found. When running in a Github Actions Context, provide the " From d703b13aee661af833485f638587458113346d5a Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 1 Apr 2025 14:44:17 -0700 Subject: [PATCH 035/122] replay tests are now grouped by benchmark file. each benchmark test file will create one replay test file. --- .../test_benchmark_bubble_sort.py | 6 ++ codeflash/benchmarking/replay_test.py | 89 ++++++++++--------- codeflash/models/models.py | 2 +- tests/test_trace_benchmarks.py | 47 ++++++---- 4 files changed, 86 insertions(+), 58 deletions(-) diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py index 03b9d38d1..21f9755a5 100644 --- a/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py @@ -15,6 +15,12 @@ def test_sort2(): def test_class_sort(benchmark): obj = Sorter(list(reversed(range(100)))) result1 = benchmark(obj.sorter, 2) + +def test_class_sort2(benchmark): result2 = benchmark(Sorter.sort_class, list(reversed(range(100)))) + +def test_class_sort3(benchmark): result3 = benchmark(Sorter.sort_static, list(reversed(range(100)))) + +def test_class_sort4(benchmark): result4 = benchmark(Sorter, [1,2,3]) \ No newline at end of file diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 9ecac2ec4..6466b24db 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -16,7 +16,7 @@ def get_next_arg_and_return( - trace_file: str, function_name: str, file_path: str, class_name: str | None = None, num_to_get: int = 256 + trace_file: str, benchmark_function_name:str, function_name: str, file_path: str, class_name: str | None = None, num_to_get: int = 256 ) -> Generator[Any]: db = sqlite3.connect(trace_file) cur = db.cursor() @@ -24,13 +24,13 @@ def get_next_arg_and_return( if class_name is not None: cursor = cur.execute( - "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_path = ? AND class_name = ? LIMIT ?", - (function_name, file_path, class_name, limit), + "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = ? LIMIT ?", + (benchmark_function_name, function_name, file_path, class_name, limit), ) else: cursor = cur.execute( - "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_path = ? AND class_name = '' LIMIT ?", - (function_name, file_path, limit), + "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = '' LIMIT ?", + (benchmark_function_name, function_name, file_path, limit), ) while (val := cursor.fetchone()) is not None: @@ -61,6 +61,7 @@ def create_trace_replay_test_code( """ assert test_framework in ["pytest", "unittest"] + # Create Imports imports = f"""import dill as pickle {"import unittest" if test_framework == "unittest" else ""} from codeflash.benchmarking.replay_test import get_next_arg_and_return @@ -82,16 +83,15 @@ def create_trace_replay_test_code( imports += "\n".join(function_imports) - functions_to_optimize = [func.get("function_name") for func in functions_data - if func.get("function_name") != "__init__"] + functions_to_optimize = sorted({func.get("function_name") for func in functions_data + if func.get("function_name") != "__init__"}) metadata = f"""functions = {functions_to_optimize} trace_file_path = r"{trace_file}" """ - # Templates for different types of tests test_function_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = {function_name}(*args, **kwargs) @@ -100,7 +100,7 @@ def create_trace_replay_test_code( test_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} function_name = "{orig_function_name}" @@ -115,7 +115,7 @@ def create_trace_replay_test_code( test_class_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} if not args: @@ -125,13 +125,15 @@ def create_trace_replay_test_code( ) test_static_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} ret = {class_name_alias}{method_name}(*args, **kwargs) """ ) + # Create main body + if test_framework == "unittest": self = "self" test_template = "\nclass TestTracedFunctions(unittest.TestCase):\n" @@ -140,17 +142,20 @@ def create_trace_replay_test_code( self = "" for func in functions_data: + module_name = func.get("module_name") function_name = func.get("function_name") class_name = func.get("class_name") file_path = func.get("file_path") + benchmark_function_name = func.get("benchmark_function_name") function_properties = func.get("function_properties") if not class_name: alias = get_function_alias(module_name, function_name) test_body = test_function_body.format( + benchmark_function_name=benchmark_function_name, + orig_function_name=function_name, function_name=alias, file_path=file_path, - orig_function_name=function_name, max_run_count=max_run_count, ) else: @@ -162,6 +167,7 @@ def create_trace_replay_test_code( method_name = "." + function_name if function_name != "__init__" else "" if function_properties.is_classmethod: test_body = test_class_method_body.format( + benchmark_function_name=benchmark_function_name, orig_function_name=function_name, file_path=file_path, class_name_alias=class_name_alias, @@ -172,6 +178,7 @@ def create_trace_replay_test_code( ) elif function_properties.is_staticmethod: test_body = test_static_method_body.format( + benchmark_function_name=benchmark_function_name, orig_function_name=function_name, file_path=file_path, class_name_alias=class_name_alias, @@ -182,6 +189,7 @@ def create_trace_replay_test_code( ) else: test_body = test_method_body.format( + benchmark_function_name=benchmark_function_name, orig_function_name=function_name, file_path=file_path, class_name_alias=class_name_alias, @@ -217,25 +225,25 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework conn = sqlite3.connect(trace_file_path.as_posix()) cursor = conn.cursor() - # Get distinct benchmark names + # Get distinct benchmark file paths cursor.execute( - "SELECT DISTINCT benchmark_function_name, benchmark_file_path FROM benchmark_function_timings" + "SELECT DISTINCT benchmark_file_path FROM benchmark_function_timings" ) - benchmarks = cursor.fetchall() + benchmark_files = cursor.fetchall() - # Generate a test for each benchmark - for benchmark in benchmarks: - benchmark_function_name, benchmark_file_path = benchmark - # Get functions associated with this benchmark + # Generate a test for each benchmark file + for benchmark_file in benchmark_files: + benchmark_file_path = benchmark_file[0] + # Get all benchmarks and functions associated with this file path cursor.execute( - "SELECT DISTINCT function_name, class_name, module_name, file_path, benchmark_line_number FROM benchmark_function_timings " - "WHERE benchmark_function_name = ? AND benchmark_file_path = ?", - (benchmark_function_name, benchmark_file_path) + "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,) ) functions_data = [] - for func_row in cursor.fetchall(): - function_name, class_name, module_name, file_path, benchmark_line_number = func_row + for row in cursor.fetchall(): + benchmark_function_name, function_name, class_name, module_name, file_path, benchmark_line_number = row # Add this function to our list functions_data.append({ "function_name": function_name, @@ -246,16 +254,15 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework "benchmark_file_path": benchmark_file_path, "benchmark_line_number": benchmark_line_number, "function_properties": inspect_top_level_functions_or_methods( - file_name=Path(file_path), - function_or_method_name=function_name, - class_name=class_name, - ) + file_name=Path(file_path), + function_or_method_name=function_name, + class_name=class_name, + ) }) if not functions_data: - logger.info(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_path}") + logger.info(f"No benchmark test functions found in {benchmark_file_path}") continue - # Generate the test code for this benchmark test_code = create_trace_replay_test_code( trace_file=trace_file_path.as_posix(), @@ -265,17 +272,15 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework ) test_code = isort.code(test_code) - # Write to file if requested - if output_dir: - name = Path(benchmark_file_path).name.split(".")[0][5:] # remove "test_" from the name since we add it in later - output_file = get_test_file_path( - test_dir=Path(output_dir), function_name=f"{name}_{benchmark_function_name}", test_type="replay" - ) - # Write test code to file, parents = true - output_dir.mkdir(parents=True, exist_ok=True) - output_file.write_text(test_code, "utf-8") - count += 1 - logger.info(f"Replay test for benchmark `{benchmark_function_name}` in {name} written to {output_file}") + name = Path(benchmark_file_path).name.split(".")[0][5:] # remove "test_" from the name since we add it in later + output_file = get_test_file_path( + test_dir=Path(output_dir), function_name=f"{name}", test_type="replay" + ) + # Write test code to file, parents = true + output_dir.mkdir(parents=True, exist_ok=True) + output_file.write_text(test_code, "utf-8") + count += 1 + logger.info(f"Replay test for benchmark file `{benchmark_file_path}` in {name} written to {output_file}") conn.close() diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 8a353b6a4..026a73a72 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -483,7 +483,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:]}_{benchmark_key.function_name}__replay_test_", project_root) + 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) 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(): diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index c67c7e87d..7851230d9 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -37,21 +37,21 @@ def test_trace_benchmarks(): process_and_bubble_sort_path = (project_root / "process_and_bubble_sort_codeflash_trace.py").as_posix() # Expected function calls expected_calls = [ - ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 20), + "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 17), ("sort_class", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 18), + "test_class_sort2", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 20), ("sort_static", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 19), + "test_class_sort3", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 23), - ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 17), + "test_class_sort4", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 26), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", @@ -76,20 +76,28 @@ def test_trace_benchmarks(): # Close connection conn.close() generate_replay_test(output_file, tests_root) - test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort_test_class_sort__replay_test_0.py") + test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort__replay_test_0.py") assert test_class_sort_path.exists() test_class_sort_code = f""" import dill as pickle from code_to_optimize.bubble_sort_codeflash_trace import \\ Sorter as code_to_optimize_bubble_sort_codeflash_trace_Sorter +from code_to_optimize.bubble_sort_codeflash_trace import \\ + sorter as code_to_optimize_bubble_sort_codeflash_trace_sorter from codeflash.benchmarking.replay_test import get_next_arg_and_return -functions = ['sorter', 'sort_class', 'sort_static'] +functions = ['sort_class', 'sort_static', 'sorter'] trace_file_path = r"{output_file.as_posix()}" +def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_sort", function_name="sorter", file_path=r"{bubble_sort_path}", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + ret = code_to_optimize_bubble_sort_codeflash_trace_sorter(*args, **kwargs) + def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_class_sort", function_name="sorter", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "sorter" @@ -102,7 +110,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): ret = instance.sorter(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_class_sort2", function_name="sort_class", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) if not args: @@ -110,13 +118,13 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_class(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_static(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_class_sort3", function_name="sort_static", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_static(*args, **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_class_sort4", function_name="__init__", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "__init__" @@ -131,20 +139,29 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): """ assert test_class_sort_path.read_text("utf-8").strip()==test_class_sort_code.strip() - test_sort_path = tests_root / Path("test_benchmark_bubble_sort_test_sort__replay_test_0.py") + test_sort_path = tests_root / Path("test_process_and_sort__replay_test_0.py") assert test_sort_path.exists() test_sort_code = f""" import dill as pickle from code_to_optimize.bubble_sort_codeflash_trace import \\ sorter as code_to_optimize_bubble_sort_codeflash_trace_sorter +from code_to_optimize.process_and_bubble_sort_codeflash_trace import \\ + compute_and_sort as \\ + code_to_optimize_process_and_bubble_sort_codeflash_trace_compute_and_sort from codeflash.benchmarking.replay_test import get_next_arg_and_return -functions = ['sorter'] +functions = ['compute_and_sort', 'sorter'] trace_file_path = r"{output_file}" +def test_code_to_optimize_process_and_bubble_sort_codeflash_trace_compute_and_sort(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_compute_and_sort", function_name="compute_and_sort", file_path=r"{process_and_bubble_sort_path}", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + ret = code_to_optimize_process_and_bubble_sort_codeflash_trace_compute_and_sort(*args, **kwargs) + def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_path=r"{bubble_sort_path}", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_no_func", function_name="sorter", file_path=r"{bubble_sort_path}", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_sorter(*args, **kwargs) From 30ec0c48713919c13283215a1c4c86a118376461 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 13:45:38 -0700 Subject: [PATCH 036/122] Use module path instead of file path for benchmarks, improved display to console. --- ... => test_benchmark_bubble_sort_example.py} | 0 ...rt.py => test_process_and_sort_example.py} | 0 codeflash/benchmarking/plugin/plugin.py | 13 ++++-- .../pytest_new_process_trace_benchmarks.py | 4 +- codeflash/benchmarking/replay_test.py | 8 +--- codeflash/benchmarking/trace_benchmarks.py | 9 ++-- codeflash/benchmarking/utils.py | 6 ++- codeflash/models/models.py | 1 - codeflash/optimization/function_optimizer.py | 38 ++++++--------- codeflash/optimization/optimizer.py | 24 ++++------ codeflash/result/explanation.py | 46 +++++++++++++++++-- tests/test_trace_benchmarks.py | 38 +++++++-------- 12 files changed, 106 insertions(+), 81 deletions(-) rename code_to_optimize/tests/pytest/benchmarks_test/{test_benchmark_bubble_sort.py => test_benchmark_bubble_sort_example.py} (100%) rename code_to_optimize/tests/pytest/benchmarks_test/{test_process_and_sort.py => test_process_and_sort_example.py} (100%) diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort_example.py similarity index 100% rename from code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py rename to code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort_example.py diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort_example.py similarity index 100% rename from code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py rename to code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort_example.py diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index b022f9afb..fc19b19d5 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -1,11 +1,15 @@ from __future__ import annotations + import os import sqlite3 import sys import time from pathlib import Path + import pytest + from codeflash.benchmarking.codeflash_trace import codeflash_trace +from codeflash.code_utils.code_utils import module_name_from_file_path from codeflash.models.models import BenchmarkKey @@ -13,11 +17,13 @@ class CodeFlashBenchmarkPlugin: def __init__(self) -> None: self._trace_path = None self._connection = None + self.project_root = None self.benchmark_timings = [] - def setup(self, trace_path:str) -> None: + def setup(self, trace_path:str, project_root:str) -> None: try: # Open connection + self.project_root = project_root self._trace_path = trace_path self._connection = sqlite3.connect(self._trace_path) cur = self._connection.cursor() @@ -235,9 +241,10 @@ def test_something(benchmark): Returns: The return value of the function + a """ - benchmark_file_path = str(self.request.node.fspath) + benchmark_file_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 @@ -273,4 +280,4 @@ def benchmark(request): return CodeFlashBenchmarkPlugin.Benchmark(request) -codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin() \ No newline at end of file +codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 7b6bd747a..1bb7bbfa4 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -13,7 +13,7 @@ import pytest try: - codeflash_benchmark_plugin.setup(trace_file) + 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] @@ -22,4 +22,4 @@ except Exception as e: print(f"Failed to collect tests: {e!s}", file=sys.stderr) exitcode = -1 - sys.exit(exitcode) \ No newline at end of file + sys.exit(exitcode) diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 6466b24db..5b654de92 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -271,20 +271,16 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework max_run_count=max_run_count, ) test_code = isort.code(test_code) - - name = Path(benchmark_file_path).name.split(".")[0][5:] # remove "test_" from the name since we add it in later output_file = get_test_file_path( - test_dir=Path(output_dir), function_name=f"{name}", test_type="replay" + test_dir=Path(output_dir), function_name=benchmark_file_path, test_type="replay" ) # Write test code to file, parents = true output_dir.mkdir(parents=True, exist_ok=True) output_file.write_text(test_code, "utf-8") count += 1 - logger.info(f"Replay test for benchmark file `{benchmark_file_path}` in {name} written to {output_file}") conn.close() - except Exception as e: logger.info(f"Error generating replay tests: {e}") - return count \ No newline at end of file + return count diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 8882078d9..8f68030cb 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -1,13 +1,12 @@ from __future__ import annotations import re - -from pytest import ExitCode +import subprocess +from pathlib import Path from codeflash.cli_cmds.console import logger from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE -from pathlib import Path -import subprocess + def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, trace_file: Path, timeout:int = 300) -> None: result = subprocess.run( @@ -40,4 +39,4 @@ def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root error_section = result.stdout logger.warning( f"Error collecting benchmarks - Pytest Exit code: {result.returncode}, {error_section}" - ) \ No newline at end of file + ) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 1d8b22f50..212512ea9 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,4 +1,5 @@ from __future__ import annotations + from typing import Optional from rich.console import Console @@ -6,7 +7,7 @@ from codeflash.cli_cmds.console import logger from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.models.models import ProcessedBenchmarkInfo, BenchmarkDetail, BenchmarkKey +from codeflash.models.models import BenchmarkDetail, BenchmarkKey, ProcessedBenchmarkInfo from codeflash.result.critic import performance_gain @@ -37,6 +38,7 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, di def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey, float, float, float]]]) -> None: console = Console() for func_path, sorted_tests in function_to_results.items(): + console.print() function_name = func_path.split(":")[-1] # Create a table for this function @@ -114,4 +116,4 @@ def process_benchmark_data( ) ) - return ProcessedBenchmarkInfo(benchmark_details=benchmark_details) \ No newline at end of file + return ProcessedBenchmarkInfo(benchmark_details=benchmark_details) diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 026a73a72..3e4848998 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -479,7 +479,6 @@ def merge(self, other: TestResults) -> None: def group_by_benchmarks(self, benchmark_keys:list[BenchmarkKey], benchmark_replay_test_dir: Path, project_root: Path) -> dict[BenchmarkKey, TestResults]: """Group TestResults by benchmark for calculating improvements for each benchmark.""" - test_results_by_benchmark = defaultdict(TestResults) benchmark_module_path = {} for benchmark_key in benchmark_keys: diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 807fd3a8c..53a342057 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -42,7 +42,6 @@ from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast from codeflash.code_utils.time_utils import humanize_runtime from codeflash.context import code_context_extractor -from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.either import Failure, Success, is_successful from codeflash.models.ExperimentMetadata import ExperimentMetadata from codeflash.models.models import ( @@ -58,7 +57,7 @@ TestFiles, TestingMode, TestResults, - TestType, BenchmarkKey, + TestType, ) from codeflash.result.create_pr import check_create_pr, existing_tests_source_for from codeflash.result.critic import coverage_critic, performance_gain, quantity_of_tests_critic, speedup_critic @@ -75,8 +74,9 @@ if TYPE_CHECKING: from argparse import Namespace + from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.either import Result - from codeflash.models.models import CoverageData, FunctionSource, OptimizedCandidate + from codeflash.models.models import BenchmarkKey, CoverageData, FunctionSource, OptimizedCandidate from codeflash.verification.verification_utils import TestConfig @@ -92,6 +92,7 @@ def __init__( function_benchmark_timings: dict[BenchmarkKey, int] | None = None, total_benchmark_timings: dict[BenchmarkKey, int] | None = None, args: Namespace | None = None, + replay_tests_dir: Path|None = None ) -> None: self.project_root = test_cfg.project_root_path self.test_cfg = test_cfg @@ -120,6 +121,7 @@ def __init__( self.function_benchmark_timings = function_benchmark_timings if function_benchmark_timings else {} self.total_benchmark_timings = total_benchmark_timings if total_benchmark_timings else {} + self.replay_tests_dir = replay_tests_dir if replay_tests_dir else None def optimize_function(self) -> Result[BestOptimization, str]: should_run_experiment = self.experiment_id is not None @@ -392,7 +394,7 @@ def determine_best_candidate( ) continue - # Instrument codeflash capture + run_results = self.run_optimized_candidate( optimization_candidate_index=candidate_index, baseline_results=original_code_baseline, @@ -430,8 +432,8 @@ def determine_best_candidate( tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") replay_perf_gain = {} if self.args.benchmark: - logger.info(f"Calculating benchmark improvement..") - test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) + benchmark_tree = Tree("Speedup percentage on benchmarks:") + test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.replay_tests_dir, self.project_root) for benchmark_key, candidate_test_results in test_results_by_benchmark.items(): original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results[benchmark_key].total_passed_runtime() candidate_replay_runtime = candidate_test_results.total_passed_runtime() @@ -439,15 +441,8 @@ def determine_best_candidate( original_runtime_ns=original_code_replay_runtime, optimized_runtime_ns=candidate_replay_runtime, ) - tree.add( - f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") - tree.add( - f"Best benchmark replay runtime: {humanize_runtime(candidate_replay_runtime)} " - f"(measured over {candidate_result.max_loop_count} " - f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" - ) - tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain[benchmark_key] * 100:.1f}%") - tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain[benchmark_key] + 1:.1f}X") + benchmark_tree.add(f"{benchmark_key}: {replay_perf_gain[benchmark_key] * 100:.1f}%") + best_optimization = BestOptimization( candidate=candidate, helper_functions=code_context.helper_functions, @@ -467,6 +462,8 @@ def determine_best_candidate( tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X") console.print(tree) + if self.args.benchmark and benchmark_tree: + console.print(benchmark_tree) console.rule() self.write_code_and_helpers( @@ -903,10 +900,7 @@ def establish_original_code_baseline( logger.debug(f"Total original code runtime (ns): {total_timing}") if self.args.benchmark: - replay_benchmarking_test_results = benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) - for benchmark_name, benchmark_results in replay_benchmarking_test_results.items(): - - logger.info(f"Replay benchmark '{benchmark_name}' runtime: {humanize_runtime(benchmark_results.total_passed_runtime())}") + replay_benchmarking_test_results = benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.replay_tests_dir, self.project_root) return Success( ( OriginalCodeBaseline( @@ -929,7 +923,6 @@ def run_optimized_candidate( file_path_to_helper_classes: dict[Path, set[str]], ) -> Result[OptimizedCandidateResult, str]: assert (test_framework := self.args.test_framework) in ["pytest", "unittest"] - with progress_bar("Testing optimization candidate"): test_env = os.environ.copy() test_env["CODEFLASH_LOOP_INDEX"] = "0" @@ -941,8 +934,6 @@ def run_optimized_candidate( test_env["PYTHONPATH"] += os.pathsep + str(self.project_root) get_run_tmp_file(Path(f"test_return_values_{optimization_candidate_index}.sqlite")).unlink(missing_ok=True) - get_run_tmp_file(Path(f"test_return_values_{optimization_candidate_index}.sqlite")).unlink(missing_ok=True) - # Instrument codeflash capture candidate_fto_code = Path(self.function_to_optimize.file_path).read_text("utf-8") candidate_helper_code = {} @@ -973,7 +964,6 @@ def run_optimized_candidate( ) ) console.rule() - if compare_test_results(baseline_results.behavioral_test_results, candidate_behavior_results): logger.info("Test results matched!") console.rule() @@ -1027,7 +1017,7 @@ def run_optimized_candidate( logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") if self.args.benchmark: - candidate_replay_benchmarking_results = candidate_benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) + candidate_replay_benchmarking_results = candidate_benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.replay_tests_dir, self.project_root) for benchmark_name, benchmark_results in candidate_replay_benchmarking_results.items(): logger.debug(f"Benchmark {benchmark_name} runtime (ns): {humanize_runtime(benchmark_results.total_passed_runtime())}") return Success( diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 219a90a79..5e162f513 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -4,10 +4,12 @@ import os import shutil import tempfile +from collections import defaultdict from pathlib import Path from typing import TYPE_CHECKING from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient +from codeflash.benchmarking.instrument_codeflash_trace import instrument_codeflash_trace_decorator from codeflash.benchmarking.plugin.plugin import CodeFlashBenchmarkPlugin from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest @@ -20,16 +22,10 @@ from codeflash.discovery.discover_unit_tests import discover_unit_tests from codeflash.discovery.functions_to_optimize import get_functions_to_optimize from codeflash.either import is_successful -from codeflash.models.models import TestType, ValidCode -from codeflash.models.models import ValidCode, TestType, BenchmarkKey +from codeflash.models.models import BenchmarkKey, TestType, ValidCode from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.telemetry.posthog_cf import ph from codeflash.verification.verification_utils import TestConfig -from codeflash.benchmarking.utils import print_benchmark_table -from codeflash.benchmarking.instrument_codeflash_trace import instrument_codeflash_trace_decorator - - -from collections import defaultdict if TYPE_CHECKING: from argparse import Namespace @@ -54,7 +50,7 @@ def __init__(self, args: Namespace) -> None: self.aiservice_client = AiServiceClient() self.experiment_id = os.getenv("CODEFLASH_EXPERIMENT_ID", None) self.local_aiservice_client = LocalAiServiceClient() if self.experiment_id else None - + self.replay_tests_dir = None def create_function_optimizer( self, function_to_optimize: FunctionToOptimize, @@ -74,6 +70,7 @@ def create_function_optimizer( args=self.args, function_benchmark_timings=function_benchmark_timings if function_benchmark_timings else None, total_benchmark_timings=total_benchmark_timings if total_benchmark_timings else None, + replay_tests_dir = self.replay_tests_dir ) def run(self) -> None: @@ -115,9 +112,9 @@ def run(self) -> None: if trace_file.exists(): trace_file.unlink() - replay_tests_dir = Path(self.args.benchmarks_root) / "codeflash_replay_tests" + self.replay_tests_dir = Path(tempfile.mkdtemp(prefix="codeflash_replay_tests_", dir=self.args.benchmarks_root)) trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Run all tests that use pytest-benchmark - replay_count = generate_replay_test(trace_file, replay_tests_dir) + replay_count = generate_replay_test(trace_file, self.replay_tests_dir) if replay_count == 0: logger.info(f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization") else: @@ -125,10 +122,9 @@ def run(self) -> None: total_benchmark_timings = CodeFlashBenchmarkPlugin.get_benchmark_timings(trace_file) function_to_results = validate_and_format_benchmark_table(function_benchmark_timings, total_benchmark_timings) print_benchmark_table(function_to_results) - logger.info("Finished tracing existing benchmarks") except Exception as e: logger.info(f"Error while tracing existing benchmarks: {e}") - logger.info(f"Information on existing benchmarks will not be available for this run.") + logger.info("Information on existing benchmarks will not be available for this run.") finally: # Restore original source code for file in file_path_to_source_code: @@ -250,8 +246,8 @@ def run(self) -> None: if function_optimizer.test_cfg.concolic_test_root_dir: shutil.rmtree(function_optimizer.test_cfg.concolic_test_root_dir, ignore_errors=True) if self.args.benchmark: - # if replay_tests_dir.exists(): - # shutil.rmtree(replay_tests_dir, ignore_errors=True) + if self.replay_tests_dir.exists(): + shutil.rmtree(self.replay_tests_dir, ignore_errors=True) trace_file.unlink(missing_ok=True) if hasattr(get_run_tmp_file, "tmpdir"): get_run_tmp_file.tmpdir.cleanup() diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index e56558a94..75288bb60 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -1,8 +1,13 @@ from __future__ import annotations + +import shutil +from io import StringIO from pathlib import Path -from typing import Optional, Union +from typing import Optional from pydantic.dataclasses import dataclass +from rich.console import Console +from rich.table import Table from codeflash.code_utils.time_utils import humanize_runtime from codeflash.models.models import BenchmarkDetail, TestResults @@ -43,11 +48,42 @@ def to_console_string(self) -> str: benchmark_info = "" if self.benchmark_details: - benchmark_info += "Benchmark Performance Details:\n" + # Get terminal width (or use a reasonable default if detection fails) + try: + terminal_width = int(shutil.get_terminal_size().columns * 0.8) + except Exception: + terminal_width = 200 # Fallback width + + # Create a rich table for better formatting + table = Table(title="Benchmark Performance Details", width=terminal_width) + + # Add columns - split Benchmark File and Function into separate columns + # Using proportional width for benchmark file column (40% of terminal width) + benchmark_col_width = max(int(terminal_width * 0.4), 40) + table.add_column("Benchmark File", style="cyan", width=benchmark_col_width) + table.add_column("Function", style="cyan") + table.add_column("Original Runtime", style="magenta") + table.add_column("Expected New Runtime", style="green") + table.add_column("Speedup", style="red") + + # Add rows with split data for detail in self.benchmark_details: - benchmark_info += f"Original timing for {detail.benchmark_name}::{detail.test_function}: {detail.original_timing}\n" - benchmark_info += f"Expected new timing for {detail.benchmark_name}::{detail.test_function}: {detail.expected_new_timing}\n" - benchmark_info += f"Benchmark speedup for {detail.benchmark_name}::{detail.test_function}: {detail.speedup_percent:.2f}%\n\n" + # Split the benchmark name and test function + benchmark_name = detail.benchmark_name + test_function = detail.test_function + + table.add_row( + benchmark_name, + test_function, + f"{detail.original_timing}", + f"{detail.expected_new_timing}", + f"{detail.speedup_percent:.2f}%" + ) + + # Render table to string - using actual terminal width + console = Console(file=StringIO(), width=terminal_width) + console.print(table) + benchmark_info = console.file.getvalue() + "\n" return ( f"Optimized {self.function_name} in {self.file_path}\n" diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 7851230d9..c50ab4332 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -13,9 +13,9 @@ 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_test" - tests_root = project_root / "tests" / "test_trace_benchmarks" - tests_root.mkdir(parents=False, exist_ok=False) - output_file = (tests_root / Path("test_trace_benchmarks.trace")).resolve() + replay_tests_dir = benchmarks_root / "codeflash_replay_tests" + 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: @@ -39,31 +39,31 @@ def test_trace_benchmarks(): expected_calls = [ ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 17), + "test_class_sort", "tests.pytest.benchmarks_test.test_benchmark_bubble_sort_example", 17), ("sort_class", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort2", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 20), + "test_class_sort2", "tests.pytest.benchmarks_test.test_benchmark_bubble_sort_example", 20), ("sort_static", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort3", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 23), + "test_class_sort3", "tests.pytest.benchmarks_test.test_benchmark_bubble_sort_example", 23), ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort4", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 26), + "test_class_sort4", "tests.pytest.benchmarks_test.test_benchmark_bubble_sort_example", 26), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 7), + "test_sort", "tests.pytest.benchmarks_test.test_benchmark_bubble_sort_example", 7), ("compute_and_sort", "", "code_to_optimize.process_and_bubble_sort_codeflash_trace", f"{process_and_bubble_sort_path}", - "test_compute_and_sort", str(benchmarks_root / "test_process_and_sort.py"), 4), + "test_compute_and_sort", "tests.pytest.benchmarks_test.test_process_and_sort_example", 4), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_no_func", str(benchmarks_root / "test_process_and_sort.py"), 8), + "test_no_func", "tests.pytest.benchmarks_test.test_process_and_sort_example", 8), ] for idx, (actual, expected) in enumerate(zip(function_calls, expected_calls)): assert actual[0] == expected[0], f"Mismatch at index {idx} for function_name" @@ -75,8 +75,8 @@ def test_trace_benchmarks(): assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number" # Close connection conn.close() - generate_replay_test(output_file, tests_root) - test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort__replay_test_0.py") + generate_replay_test(output_file, replay_tests_dir) + test_class_sort_path = replay_tests_dir/ Path("test_tests_pytest_benchmarks_test_test_benchmark_bubble_sort_example__replay_test_0.py") assert test_class_sort_path.exists() test_class_sort_code = f""" import dill as pickle @@ -139,7 +139,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): """ assert test_class_sort_path.read_text("utf-8").strip()==test_class_sort_code.strip() - test_sort_path = tests_root / Path("test_process_and_sort__replay_test_0.py") + test_sort_path = replay_tests_dir / Path("test_tests_pytest_benchmarks_test_test_process_and_sort_example__replay_test_0.py") assert test_sort_path.exists() test_sort_code = f""" import dill as pickle @@ -170,14 +170,14 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): assert test_sort_path.read_text("utf-8").strip()==test_sort_code.strip() finally: # cleanup - shutil.rmtree(tests_root) + output_file.unlink(missing_ok=True) + shutil.rmtree(replay_tests_dir) def test_trace_multithreaded_benchmark() -> None: project_root = Path(__file__).parent.parent / "code_to_optimize" benchmarks_root = project_root / "tests" / "pytest" / "benchmarks_multithread" - tests_root = project_root / "tests" / "test_trace_benchmarks" - tests_root.mkdir(parents=False, exist_ok=False) - output_file = (tests_root / Path("test_trace_benchmarks.trace")).resolve() + 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: @@ -209,7 +209,7 @@ def test_trace_multithreaded_benchmark() -> None: expected_calls = [ ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_benchmark_sort", str(benchmarks_root / "test_multithread_sort.py"), 4), + "test_benchmark_sort", "tests.pytest.benchmarks_multithread.test_multithread_sort", 4), ] for idx, (actual, expected) in enumerate(zip(function_calls, expected_calls)): assert actual[0] == expected[0], f"Mismatch at index {idx} for function_name" @@ -224,4 +224,4 @@ def test_trace_multithreaded_benchmark() -> None: finally: # cleanup - shutil.rmtree(tests_root) \ No newline at end of file + output_file.unlink(missing_ok=True) \ No newline at end of file From bb9c5db4b10aa463a6d6608cb2fd1da218b1fa56 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 15:12:19 -0700 Subject: [PATCH 037/122] benchmark flow is working. changed paths to use module_path instead of file_path for Benchmarkkey --- codeflash/benchmarking/codeflash_trace.py | 17 +-- codeflash/benchmarking/plugin/plugin.py | 24 ++-- codeflash/benchmarking/replay_test.py | 14 +-- codeflash/benchmarking/utils.py | 26 +++- codeflash/models/models.py | 6 +- codeflash/optimization/function_optimizer.py | 118 ++++++++++--------- codeflash/result/explanation.py | 4 +- tests/test_trace_benchmarks.py | 8 +- 8 files changed, 117 insertions(+), 100 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index bcbb3268c..a2d080283 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/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 diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index fc19b19d5..c7c11c6d4 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/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 diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 5b654de92..63a330774 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/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) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 212512ea9..dff32b57e 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,5 +1,6 @@ from __future__ import annotations +import shutil from typing import Optional from rich.console import Console @@ -35,8 +36,14 @@ 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] @@ -44,23 +51,30 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey # 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)), diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 01139e3d3..a586b8d6f 100644 --- a/codeflash/models/models.py +++ b/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(): diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 6fc771b29..da9378f12 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -426,13 +426,13 @@ def determine_best_candidate( continue - run_results = self.run_optimized_candidate( - optimization_candidate_index=candidate_index, - baseline_results=original_code_baseline, - original_helper_code=original_helper_code, - file_path_to_helper_classes=file_path_to_helper_classes, - ) - console.rule() + run_results = self.run_optimized_candidate( + optimization_candidate_index=candidate_index, + baseline_results=original_code_baseline, + original_helper_code=original_helper_code, + file_path_to_helper_classes=file_path_to_helper_classes, + ) + console.rule() if not is_successful(run_results): optimized_runtimes[candidate.optimization_id] = None @@ -448,58 +448,60 @@ def determine_best_candidate( ) speedup_ratios[candidate.optimization_id] = perf_gain - tree = Tree(f"Candidate #{candidate_index} - Runtime Information") - if speedup_critic( - candidate_result, original_code_baseline.runtime, best_runtime_until_now - ) and quantity_of_tests_critic(candidate_result): - tree.add("This candidate is faster than the previous best candidate. πŸš€") - tree.add(f"Original summed runtime: {humanize_runtime(original_code_baseline.runtime)}") - tree.add( - f"Best summed runtime: {humanize_runtime(candidate_result.best_test_runtime)} " - f"(measured over {candidate_result.max_loop_count} " - f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" - ) - tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") - replay_perf_gain = {} - if self.args.benchmark: - benchmark_tree = Tree("Speedup percentage on benchmarks:") - test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.replay_tests_dir, self.project_root) - for benchmark_key, candidate_test_results in test_results_by_benchmark.items(): - original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results[benchmark_key].total_passed_runtime() - candidate_replay_runtime = candidate_test_results.total_passed_runtime() - replay_perf_gain[benchmark_key] = performance_gain( - original_runtime_ns=original_code_replay_runtime, - optimized_runtime_ns=candidate_replay_runtime, - ) - benchmark_tree.add(f"{benchmark_key}: {replay_perf_gain[benchmark_key] * 100:.1f}%") - - best_optimization = BestOptimization( - candidate=candidate, - helper_functions=code_context.helper_functions, - runtime=best_test_runtime, - winning_behavioral_test_results=candidate_result.behavior_test_results, - replay_performance_gain=replay_perf_gain if self.args.benchmark else None, - winning_benchmarking_test_results=candidate_result.benchmarking_test_results, - winning_replay_benchmarking_test_results=candidate_result.benchmarking_test_results, - ) - best_runtime_until_now = best_test_runtime - else: - tree.add( - f"Summed runtime: {humanize_runtime(best_test_runtime)} " - f"(measured over {candidate_result.max_loop_count} " - f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" - ) - tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X") - console.print(tree) - if self.args.benchmark and benchmark_tree: - console.print(benchmark_tree) - console.rule() + tree = Tree(f"Candidate #{candidate_index} - Runtime Information") + if speedup_critic( + candidate_result, original_code_baseline.runtime, best_runtime_until_now + ) and quantity_of_tests_critic(candidate_result): + tree.add("This candidate is faster than the previous best candidate. πŸš€") + tree.add(f"Original summed runtime: {humanize_runtime(original_code_baseline.runtime)}") + tree.add( + f"Best summed runtime: {humanize_runtime(candidate_result.best_test_runtime)} " + f"(measured over {candidate_result.max_loop_count} " + f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + ) + tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") + replay_perf_gain = {} + if self.args.benchmark: + test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.replay_tests_dir, self.project_root) + if len(test_results_by_benchmark) > 0: + benchmark_tree = Tree("Speedup percentage on benchmarks:") + for benchmark_key, candidate_test_results in test_results_by_benchmark.items(): + + original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results[benchmark_key].total_passed_runtime() + candidate_replay_runtime = candidate_test_results.total_passed_runtime() + replay_perf_gain[benchmark_key] = performance_gain( + original_runtime_ns=original_code_replay_runtime, + optimized_runtime_ns=candidate_replay_runtime, + ) + benchmark_tree.add(f"{benchmark_key}: {replay_perf_gain[benchmark_key] * 100:.1f}%") + + best_optimization = BestOptimization( + candidate=candidate, + helper_functions=code_context.helper_functions, + runtime=best_test_runtime, + winning_behavioral_test_results=candidate_result.behavior_test_results, + replay_performance_gain=replay_perf_gain if self.args.benchmark else None, + winning_benchmarking_test_results=candidate_result.benchmarking_test_results, + winning_replay_benchmarking_test_results=candidate_result.benchmarking_test_results, + ) + best_runtime_until_now = best_test_runtime + else: + tree.add( + f"Summed runtime: {humanize_runtime(best_test_runtime)} " + f"(measured over {candidate_result.max_loop_count} " + f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + ) + tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X") + console.print(tree) + if self.args.benchmark and benchmark_tree: + console.print(benchmark_tree) + console.rule() - self.write_code_and_helpers( - self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path - ) + self.write_code_and_helpers( + self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path + ) except KeyboardInterrupt as e: self.write_code_and_helpers( diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 75288bb60..076fb0e55 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -60,8 +60,8 @@ def to_console_string(self) -> str: # Add columns - split Benchmark File and Function into separate columns # Using proportional width for benchmark file column (40% of terminal width) benchmark_col_width = max(int(terminal_width * 0.4), 40) - table.add_column("Benchmark File", style="cyan", width=benchmark_col_width) - table.add_column("Function", style="cyan") + table.add_column("Benchmark Module Path", style="cyan", width=benchmark_col_width) + table.add_column("Test Function", style="cyan") table.add_column("Original Runtime", style="magenta") table.add_column("Expected New Runtime", style="green") table.add_column("Speedup", style="red") diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index c50ab4332..e953d1e81 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -27,7 +27,7 @@ def test_trace_benchmarks(): # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_path, benchmark_function_name, benchmark_file_path, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_path, benchmark_function_name, function_name") + "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 @@ -71,7 +71,7 @@ def test_trace_benchmarks(): 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_file_path" + 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() @@ -189,7 +189,7 @@ def test_trace_multithreaded_benchmark() -> None: # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_path, benchmark_function_name, benchmark_file_path, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_path, benchmark_function_name, function_name") + "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 @@ -217,7 +217,7 @@ def test_trace_multithreaded_benchmark() -> None: 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_file_path" + 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() From 1928dc4f6cd1744fdda24ff92efa5727e28b536b Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 15:17:08 -0700 Subject: [PATCH 038/122] fixed string error --- codeflash/models/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/models/models.py b/codeflash/models/models.py index a586b8d6f..74c963a32 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -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_{benchmark_key.module_path.replace(".", "_")}__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(): From 217e239a60915702f2aaddb2dee4c452a0bc41bb Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 15:36:43 -0700 Subject: [PATCH 039/122] fixed mypy error --- codeflash/result/explanation.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 076fb0e55..217010af2 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -3,7 +3,7 @@ import shutil from io import StringIO from pathlib import Path -from typing import Optional +from typing import Optional, cast from pydantic.dataclasses import dataclass from rich.console import Console @@ -79,11 +79,11 @@ def to_console_string(self) -> str: f"{detail.expected_new_timing}", f"{detail.speedup_percent:.2f}%" ) - - # Render table to string - using actual terminal width - console = Console(file=StringIO(), width=terminal_width) + # Convert table to string + string_buffer = StringIO() + console = Console(file=string_buffer, width=terminal_width) console.print(table) - benchmark_info = console.file.getvalue() + "\n" + benchmark_info = cast(StringIO, console.file).getvalue() + "\n" # Cast for mypy return ( f"Optimized {self.function_name} in {self.file_path}\n" From 96dd78092ec0a4f3a514404474dcaabd5f7323f4 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:15:27 -0700 Subject: [PATCH 040/122] new end to end test for benchmarking bubble sort --- .../workflows/end-to-end-benchmark-test.yaml | 41 +++++++++++++++++++ codeflash/optimization/function_optimizer.py | 7 ++-- .../scripts/end_to_end_test_benchmark_sort.py | 26 ++++++++++++ tests/scripts/end_to_end_test_utilities.py | 8 ++-- 4 files changed, 76 insertions(+), 6 deletions(-) create mode 100644 .github/workflows/end-to-end-benchmark-test.yaml create mode 100644 tests/scripts/end_to_end_test_benchmark_sort.py diff --git a/.github/workflows/end-to-end-benchmark-test.yaml b/.github/workflows/end-to-end-benchmark-test.yaml new file mode 100644 index 000000000..efdb5764f --- /dev/null +++ b/.github/workflows/end-to-end-benchmark-test.yaml @@ -0,0 +1,41 @@ +name: end-to-end-test + +on: + pull_request: + workflow_dispatch: + +jobs: + benchmark-bubble-sort-optimization: + runs-on: ubuntu-latest + env: + CODEFLASH_AIS_SERVER: prod + POSTHOG_API_KEY: ${{ secrets.POSTHOG_API_KEY }} + CODEFLASH_API_KEY: ${{ secrets.CODEFLASH_API_KEY }} + COLUMNS: 110 + MAX_RETRIES: 3 + RETRY_DELAY: 5 + EXPECTED_IMPROVEMENT_PCT: 5 + CODEFLASH_END_TO_END: 1 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + token: ${{ secrets.GITHUB_TOKEN }} + + - name: Set up Python 3.11 for CLI + uses: astral-sh/setup-uv@v5 + with: + python-version: 3.11.6 + + - name: Install dependencies (CLI) + run: | + uv tool install poetry + uv venv + source .venv/bin/activate + poetry install --with dev + + - name: Run Codeflash to optimize code + id: optimize_code with benchmarks + run: | + source .venv/bin/activate + poetry run python tests/scripts/end_to_end_test_benchmark_sort.py \ No newline at end of file diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index d77976a47..1cc854736 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -449,6 +449,7 @@ def determine_best_candidate( speedup_ratios[candidate.optimization_id] = perf_gain tree = Tree(f"Candidate #{candidate_index} - Runtime Information") + benchmark_tree = None if speedup_critic( candidate_result, original_code_baseline.runtime, best_runtime_until_now ) and quantity_of_tests_critic(candidate_result): @@ -499,9 +500,9 @@ def determine_best_candidate( console.print(benchmark_tree) console.rule() - self.write_code_and_helpers( - self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path - ) + self.write_code_and_helpers( + self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path + ) except KeyboardInterrupt as e: self.write_code_and_helpers( diff --git a/tests/scripts/end_to_end_test_benchmark_sort.py b/tests/scripts/end_to_end_test_benchmark_sort.py new file mode 100644 index 000000000..64aabe384 --- /dev/null +++ b/tests/scripts/end_to_end_test_benchmark_sort.py @@ -0,0 +1,26 @@ +import os +import pathlib + +from end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command, run_with_retries + + +def run_test(expected_improvement_pct: int) -> bool: + cwd = (pathlib.Path(__file__).parent.parent.parent / "code_to_optimize").resolve() + config = TestConfig( + file_path=pathlib.Path("bubble_sort.py"), + function_name="sorter", + benchmarks_root=cwd / "tests" / "pytest" / "benchmarks", + test_framework="pytest", + min_improvement_x=1.0, + coverage_expectations=[ + CoverageExpectation( + function_name="sorter", expected_coverage=100.0, expected_lines=[2, 3, 4, 5, 6, 7, 8, 9, 10] + ) + ], + ) + + return run_codeflash_command(cwd, config, expected_improvement_pct) + + +if __name__ == "__main__": + exit(run_with_retries(run_test, int(os.getenv("EXPECTED_IMPROVEMENT_PCT", 5)))) diff --git a/tests/scripts/end_to_end_test_utilities.py b/tests/scripts/end_to_end_test_utilities.py index c961b6fd1..83ed8548c 100644 --- a/tests/scripts/end_to_end_test_utilities.py +++ b/tests/scripts/end_to_end_test_utilities.py @@ -26,6 +26,7 @@ class TestConfig: min_improvement_x: float = 0.1 trace_mode: bool = False coverage_expectations: list[CoverageExpectation] = field(default_factory=list) + benchmarks_root: Optional[pathlib.Path] = None def clear_directory(directory_path: str | pathlib.Path) -> None: @@ -85,8 +86,8 @@ def run_codeflash_command( path_to_file = cwd / config.file_path file_contents = path_to_file.read_text("utf-8") test_root = cwd / "tests" / (config.test_framework or "") - command = build_command(cwd, config, test_root) + command = build_command(cwd, config, test_root, config.benchmarks_root if config.benchmarks_root else None) process = subprocess.Popen( command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, cwd=str(cwd), env=os.environ.copy() ) @@ -116,7 +117,7 @@ def run_codeflash_command( return validated -def build_command(cwd: pathlib.Path, config: TestConfig, test_root: pathlib.Path) -> list[str]: +def build_command(cwd: pathlib.Path, config: TestConfig, test_root: pathlib.Path, benchmarks_root:pathlib.Path|None = None) -> list[str]: python_path = "../../../codeflash/main.py" if "code_directories" in str(cwd) else "../codeflash/main.py" base_command = ["python", python_path, "--file", config.file_path, "--no-pr"] @@ -127,7 +128,8 @@ def build_command(cwd: pathlib.Path, config: TestConfig, test_root: pathlib.Path base_command.extend( ["--test-framework", config.test_framework, "--tests-root", str(test_root), "--module-root", str(cwd)] ) - + if benchmarks_root: + base_command.extend(["--benchmark", "--benchmarks-root", str(benchmarks_root)]) return base_command From 5785875588579b2fd470586d0efcf6c37e2513d9 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:18:18 -0700 Subject: [PATCH 041/122] renamed test --- ...chmark-test.yaml => end-to-end-test-benchmark-bubblesort.yaml} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename .github/workflows/{end-to-end-benchmark-test.yaml => end-to-end-test-benchmark-bubblesort.yaml} (100%) diff --git a/.github/workflows/end-to-end-benchmark-test.yaml b/.github/workflows/end-to-end-test-benchmark-bubblesort.yaml similarity index 100% rename from .github/workflows/end-to-end-benchmark-test.yaml rename to .github/workflows/end-to-end-test-benchmark-bubblesort.yaml From d656d3b983b5f6c5085c3838e0d1b49731ed56d2 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:22:41 -0700 Subject: [PATCH 042/122] fixed e2e test --- .github/workflows/end-to-end-test-benchmark-bubblesort.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/end-to-end-test-benchmark-bubblesort.yaml b/.github/workflows/end-to-end-test-benchmark-bubblesort.yaml index efdb5764f..53a59dac1 100644 --- a/.github/workflows/end-to-end-test-benchmark-bubblesort.yaml +++ b/.github/workflows/end-to-end-test-benchmark-bubblesort.yaml @@ -35,7 +35,7 @@ jobs: poetry install --with dev - name: Run Codeflash to optimize code - id: optimize_code with benchmarks + id: optimize_code_with_benchmarks run: | source .venv/bin/activate poetry run python tests/scripts/end_to_end_test_benchmark_sort.py \ No newline at end of file From 4d0eb3da06996e91851f221a0d74eb8ce94c8269 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:36:47 -0700 Subject: [PATCH 043/122] printing issues on github actions --- codeflash/benchmarking/utils.py | 8 ++++---- codeflash/result/explanation.py | 12 ++++++------ 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index dff32b57e..56644167f 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -40,9 +40,9 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, di def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey, float, float, float]]]) -> None: try: - terminal_width = int(shutil.get_terminal_size().columns * 0.8) + terminal_width = int(shutil.get_terminal_size().columns * 0.9) except Exception: - terminal_width = 200 # Fallback width + terminal_width = 120 # Fallback width console = Console(width = terminal_width) for func_path, sorted_tests in function_to_results.items(): console.print() @@ -52,8 +52,8 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey table = Table(title=f"Function: {function_name}", border_style="blue") # 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("Benchmark Module Path", style="cyan", overflow="fold") + table.add_column("Test Function", style="magenta", overflow="fold") 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") diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 217010af2..2d4aba9bf 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -50,7 +50,7 @@ def to_console_string(self) -> str: if self.benchmark_details: # Get terminal width (or use a reasonable default if detection fails) try: - terminal_width = int(shutil.get_terminal_size().columns * 0.8) + terminal_width = int(shutil.get_terminal_size().columns * 0.9) except Exception: terminal_width = 200 # Fallback width @@ -60,11 +60,11 @@ def to_console_string(self) -> str: # Add columns - split Benchmark File and Function into separate columns # Using proportional width for benchmark file column (40% of terminal width) benchmark_col_width = max(int(terminal_width * 0.4), 40) - table.add_column("Benchmark Module Path", style="cyan", width=benchmark_col_width) - table.add_column("Test Function", style="cyan") - table.add_column("Original Runtime", style="magenta") - table.add_column("Expected New Runtime", style="green") - table.add_column("Speedup", style="red") + table.add_column("Benchmark Module Path", style="cyan", width=benchmark_col_width, overflow="fold") + table.add_column("Test Function", style="cyan", overflow="fold") + table.add_column("Original Runtime", style="magenta", justify="right") + table.add_column("Expected New Runtime", style="green", justify="right") + table.add_column("Speedup", style="red", justify="right") # Add rows with split data for detail in self.benchmark_details: From 6100620b5faac9bca287bb805c4ed26e33d4bd03 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:42:29 -0700 Subject: [PATCH 044/122] attempt to use horizontals for rows --- codeflash/benchmarking/utils.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 56644167f..5548feee8 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,16 +1,20 @@ from __future__ import annotations import shutil -from typing import Optional +from typing import TYPE_CHECKING, Optional +from rich.box import HORIZONTALS from rich.console import Console from rich.table import Table from codeflash.cli_cmds.console import logger from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.models.models import BenchmarkDetail, BenchmarkKey, ProcessedBenchmarkInfo +from codeflash.models.models import BenchmarkDetail, ProcessedBenchmarkInfo from codeflash.result.critic import performance_gain +if TYPE_CHECKING: + from codeflash.models.models import BenchmarkKey + def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[BenchmarkKey, int]], total_benchmark_timings: dict[BenchmarkKey, int]) -> dict[str, list[tuple[BenchmarkKey, float, float, float]]]: @@ -49,10 +53,10 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey function_name = func_path.split(":")[-1] # Create a table for this function - table = Table(title=f"Function: {function_name}", border_style="blue") - + table = Table(title=f"Function: {function_name}", width=terminal_width, border_style="blue", box=HORIZONTALS) + benchmark_col_width = max(int(terminal_width * 0.4), 40) # Add columns - split the benchmark test into two columns - table.add_column("Benchmark Module Path", style="cyan", overflow="fold") + table.add_column("Benchmark Module Path", width=benchmark_col_width, style="cyan", overflow="fold") table.add_column("Test Function", style="magenta", overflow="fold") table.add_column("Total Time (ms)", justify="right", style="green") table.add_column("Function Time (ms)", justify="right", style="yellow") From 21a79eb8c30cddbaaefe8d1210fcfd00b7bd6c8a Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:48:10 -0700 Subject: [PATCH 045/122] added row lines --- codeflash/benchmarking/utils.py | 3 +-- codeflash/result/explanation.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 5548feee8..a10b4487d 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -3,7 +3,6 @@ import shutil from typing import TYPE_CHECKING, Optional -from rich.box import HORIZONTALS from rich.console import Console from rich.table import Table @@ -53,7 +52,7 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey function_name = func_path.split(":")[-1] # Create a table for this function - table = Table(title=f"Function: {function_name}", width=terminal_width, border_style="blue", box=HORIZONTALS) + table = Table(title=f"Function: {function_name}", width=terminal_width, border_style="blue", show_lines=True) benchmark_col_width = max(int(terminal_width * 0.4), 40) # Add columns - split the benchmark test into two columns table.add_column("Benchmark Module Path", width=benchmark_col_width, style="cyan", overflow="fold") diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 2d4aba9bf..c6e1fb9dc 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -55,7 +55,7 @@ def to_console_string(self) -> str: terminal_width = 200 # Fallback width # Create a rich table for better formatting - table = Table(title="Benchmark Performance Details", width=terminal_width) + table = Table(title="Benchmark Performance Details", width=terminal_width, show_lines=True) # Add columns - split Benchmark File and Function into separate columns # Using proportional width for benchmark file column (40% of terminal width) From b374b6ea5a35479cc823359c70c59159838ef2ec Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 3 Apr 2025 15:28:44 -0700 Subject: [PATCH 046/122] made benchmarks-root use resolve() --- codeflash/cli_cmds/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 07652f707..ed0dbd760 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -135,7 +135,7 @@ def process_pyproject_config(args: Namespace) -> Namespace: if args.benchmark: assert args.benchmarks_root is not None, "--benchmarks-root must be specified when running with --benchmark" assert Path(args.benchmarks_root).is_dir(), f"--benchmarks-root {args.benchmarks_root} must be a valid directory" - assert Path(args.benchmarks_root).is_relative_to(Path(args.tests_root)), ( + assert Path(args.benchmarks_root).resolve().is_relative_to(Path(args.tests_root).resolve()), ( f"--benchmarks-root {args.benchmarks_root} must be a subdirectory of --tests-root {args.tests_root}" ) if env_utils.get_pr_number() is not None: From 27a64888485c030a3e1fc2c68eaf4fe493f43e3f Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 3 Apr 2025 15:59:32 -0700 Subject: [PATCH 047/122] handled edge case for instrumenting codeflash trace --- .../instrument_codeflash_trace.py | 15 +++- tests/test_instrument_codeflash_trace.py | 80 ++++++++++++++++++- 2 files changed, 91 insertions(+), 4 deletions(-) diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py index 06e93daf8..044b0b0a4 100644 --- a/codeflash/benchmarking/instrument_codeflash_trace.py +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -12,12 +12,14 @@ def __init__(self, target_functions: set[tuple[str, str]]) -> None: self.target_functions = target_functions self.added_codeflash_trace = False self.class_name = "" + self.function_name = "" self.decorator = cst.Decorator( decorator=cst.Name(value="codeflash_trace") ) def leave_ClassDef(self, original_node, updated_node): - self.class_name = "" + if self.class_name == original_node.name.value: + self.class_name = "" # Even if nested classes are not visited, this function is still called on them return updated_node def visit_ClassDef(self, node): @@ -25,7 +27,14 @@ def visit_ClassDef(self, node): return False self.class_name = node.name.value + def visit_FunctionDef(self, node): + if self.function_name: # Don't go into nested function + return False + self.function_name = node.name.value + def leave_FunctionDef(self, original_node, updated_node): + if self.function_name == original_node.name.value: + self.function_name = "" if (self.class_name, original_node.name.value) in self.target_functions: # Add the new decorator after any existing decorators, so it gets executed first updated_decorators = list(updated_node.decorators) + [self.decorator] @@ -33,8 +42,8 @@ def leave_FunctionDef(self, original_node, updated_node): return updated_node.with_changes( decorators=updated_decorators ) - else: - return updated_node + + return updated_node def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: # Create import statement for codeflash_trace diff --git a/tests/test_instrument_codeflash_trace.py b/tests/test_instrument_codeflash_trace.py index 6b884c631..38a6381e2 100644 --- a/tests/test_instrument_codeflash_trace.py +++ b/tests/test_instrument_codeflash_trace.py @@ -466,4 +466,82 @@ def static_method_b(): # Compare the modified content with expected content assert modified_content_1.strip() == expected_content_1.strip() - assert modified_content_2.strip() == expected_content_2.strip() \ No newline at end of file + assert modified_content_2.strip() == expected_content_2.strip() + + +def test_add_decorator_to_method_after_nested_class() -> None: + """Test adding decorator to a method that appears after a nested class definition.""" + code = """ +class OuterClass: + class NestedClass: + def nested_method(self): + return "Hello from nested class method" + + def target_method(self): + return "Hello from target method after nested class" +""" + + fto = FunctionToOptimize( + function_name="target_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="OuterClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + functions_to_optimize=[fto] + ) + + expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +class OuterClass: + class NestedClass: + def nested_method(self): + return "Hello from nested class method" + + @codeflash_trace + def target_method(self): + return "Hello from target method after nested class" +""" + + assert modified_code.strip() == expected_code.strip() + + +def test_add_decorator_to_function_after_nested_function() -> None: + """Test adding decorator to a function that appears after a function with a nested function.""" + code = """ +def function_with_nested(): + def inner_function(): + return "Hello from inner function" + + return inner_function() + +def target_function(): + return "Hello from target function after nested function" +""" + + fto = FunctionToOptimize( + function_name="target_function", + file_path=Path("dummy_path.py"), + parents=[] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + functions_to_optimize=[fto] + ) + + expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +def function_with_nested(): + def inner_function(): + return "Hello from inner function" + + return inner_function() + +@codeflash_trace +def target_function(): + return "Hello from target function after nested function" +""" + + assert modified_code.strip() == expected_code.strip() \ No newline at end of file From 4a24f2c51f9937cca0867b56abf32d98c81b42a4 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 3 Apr 2025 16:33:41 -0700 Subject: [PATCH 048/122] fixed slight bug with formatting table --- codeflash/benchmarking/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index a10b4487d..da09cd57a 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -29,7 +29,7 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, di # If the function time is greater than total time, likely to have multithreading / multiprocessing issues. # Do not try to project the optimization impact for this function. sorted_tests.append((benchmark_key, 0.0, 0.0, 0.0)) - if total_time > 0: + elif total_time > 0: percentage = (func_time / total_time) * 100 # Convert nanoseconds to milliseconds func_time_ms = func_time / 1_000_000 From 9de664bdfd67bf5323b17ef0464547854a2d0589 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 3 Apr 2025 16:48:56 -0700 Subject: [PATCH 049/122] improved file removal after errors --- codeflash/optimization/optimizer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 5e162f513..a1260dfd8 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -127,10 +127,11 @@ def run(self) -> None: logger.info("Information on existing benchmarks will not be available for this run.") finally: # Restore original source code + trace_file.unlink() + shutil.rmtree(self.replay_tests_dir, ignore_errors=True) for file in file_path_to_source_code: with file.open("w", encoding="utf8") as f: f.write(file_path_to_source_code[file]) - optimizations_found: int = 0 function_iterator_count: int = 0 if self.args.test_framework == "pytest": From a8d4fda34d4397e8226693ce860f9252b3a70056 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 4 Apr 2025 11:04:58 -0700 Subject: [PATCH 050/122] fixed a return bug --- codeflash/benchmarking/codeflash_trace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index a2d080283..776e0e635 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -144,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 None + return result if len(self.function_calls_data) > 1000: self.write_function_timings() From 1f3fcffa360c946b75547b68e43ef6c12961e477 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Mon, 7 Apr 2025 14:52:18 -0700 Subject: [PATCH 051/122] Support recursive functions, and @benchmark / @pytest.mark.benchmark ways of using benchmark. created tests for all of them --- .../bubble_sort_codeflash_trace.py | 18 +++++ .../benchmarks_test/test_recursive_example.py | 6 ++ .../test_benchmark_decorator.py | 11 +++ codeflash/benchmarking/codeflash_trace.py | 79 +++++++++++++------ codeflash/benchmarking/plugin/plugin.py | 71 ++++++++++------- .../pytest_new_process_trace_benchmarks.py | 2 +- codeflash/benchmarking/replay_test.py | 2 +- tests/test_trace_benchmarks.py | 62 ++++++++++++++- 8 files changed, 194 insertions(+), 57 deletions(-) create mode 100644 code_to_optimize/tests/pytest/benchmarks_test/test_recursive_example.py create mode 100644 code_to_optimize/tests/pytest/benchmarks_test_decorator/test_benchmark_decorator.py diff --git a/code_to_optimize/bubble_sort_codeflash_trace.py b/code_to_optimize/bubble_sort_codeflash_trace.py index ee4dbd999..48e9a412b 100644 --- a/code_to_optimize/bubble_sort_codeflash_trace.py +++ b/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): diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_recursive_example.py b/code_to_optimize/tests/pytest/benchmarks_test/test_recursive_example.py new file mode 100644 index 000000000..689b1f9ff --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks_test_decorator/test_benchmark_decorator.py b/code_to_optimize/tests/pytest/benchmarks_test_decorator/test_benchmark_decorator.py new file mode 100644 index 000000000..b924bee7f --- /dev/null +++ b/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)))) \ No newline at end of file diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 776e0e635..95318a38a 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/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, diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index c7c11c6d4..f1614b5c8 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/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,6 +211,13 @@ 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"): @@ -217,9 +225,19 @@ def pytest_collection_modifyitems(config, items): 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: @@ -227,44 +245,37 @@ 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( diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 1bb7bbfa4..232c39fa7 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/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: diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 63a330774..445957505 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/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: diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index e953d1e81..715955063 100644 --- a/tests/test_trace_benchmarks.py +++ b/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) \ No newline at end of file From fe6365262a810ab2c8ebed195d2005a3babf7e36 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 8 Apr 2025 11:32:44 -0700 Subject: [PATCH 052/122] basic pickle patch version working --- codeflash/picklepatch/__init__.py | 0 codeflash/picklepatch/pickle_patcher.py | 346 ++++++++++++++++++++ codeflash/picklepatch/pickle_placeholder.py | 66 ++++ tests/test_pickle_patcher.py | 172 ++++++++++ 4 files changed, 584 insertions(+) create mode 100644 codeflash/picklepatch/__init__.py create mode 100644 codeflash/picklepatch/pickle_patcher.py create mode 100644 codeflash/picklepatch/pickle_placeholder.py create mode 100644 tests/test_pickle_patcher.py diff --git a/codeflash/picklepatch/__init__.py b/codeflash/picklepatch/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/codeflash/picklepatch/pickle_patcher.py b/codeflash/picklepatch/pickle_patcher.py new file mode 100644 index 000000000..cfedd28fd --- /dev/null +++ b/codeflash/picklepatch/pickle_patcher.py @@ -0,0 +1,346 @@ +"""PicklePatcher - A utility for safely pickling objects with unpicklable components. + +This module provides functions to recursively pickle objects, replacing unpicklable +components with placeholders that provide informative errors when accessed. +""" + +import pickle +import types + +import dill + +from .pickle_placeholder import PicklePlaceholder + + +class PicklePatcher: + """A utility class for safely pickling objects with unpicklable components. + + This class provides methods to recursively pickle objects, replacing any + components that can't be pickled with placeholder objects. + """ + + # Class-level cache of unpicklable types + _unpicklable_types = set() + + @staticmethod + def dumps(obj, protocol=None, max_depth=100, **kwargs): + """Safely pickle an object, replacing unpicklable parts with placeholders. + + Args: + obj: The object to pickle + protocol: The pickle protocol version to use + max_depth: Maximum recursion depth + **kwargs: Additional arguments for pickle/dill.dumps + + Returns: + bytes: Pickled data with placeholders for unpicklable objects + """ + return PicklePatcher._recursive_pickle(obj, max_depth, path=[], protocol=protocol, **kwargs) + + @staticmethod + def loads(pickled_data): + """Unpickle data that may contain placeholders. + + Args: + pickled_data: Pickled data with possible placeholders + + Returns: + The unpickled object with placeholders for unpicklable parts + """ + try: + # We use dill for loading since it can handle everything pickle can + return dill.loads(pickled_data) + except Exception as e: + raise + + @staticmethod + def _create_placeholder(obj, error_msg, path): + """Create a placeholder for an unpicklable object. + + Args: + obj: The original unpicklable object + error_msg: Error message explaining why it couldn't be pickled + path: Path to this object in the object graph + + Returns: + PicklePlaceholder: A placeholder object + """ + obj_type = type(obj) + try: + obj_str = str(obj)[:100] if hasattr(obj, "__str__") else f"" + except: + obj_str = f"" + + print(f"Creating placeholder for {obj_type.__name__} at path {'->'.join(path) or 'root'}: {error_msg}") + + placeholder = PicklePlaceholder( + obj_type.__name__, + obj_str, + error_msg, + path + ) + + # Add this type to our known unpicklable types cache + PicklePatcher._unpicklable_types.add(obj_type) + return placeholder + + @staticmethod + def _pickle(obj, path=None, protocol=None, **kwargs): + """Try to pickle an object using pickle first, then dill. If both fail, create a placeholder. + + Args: + obj: The object to pickle + path: Path to this object in the object graph + protocol: The pickle protocol version to use + **kwargs: Additional arguments for pickle/dill.dumps + + Returns: + tuple: (success, result) where success is a boolean and result is either: + - Pickled bytes if successful + - Error message if not successful + """ + # Try standard pickle first + try: + return True, pickle.dumps(obj, protocol=protocol, **kwargs) + except (pickle.PickleError, TypeError, AttributeError, ValueError) as e: + # Then try dill (which is more powerful) + try: + return True, dill.dumps(obj, protocol=protocol, **kwargs) + except (dill.PicklingError, TypeError, AttributeError, ValueError) as e: + return False, str(e) + + @staticmethod + def _recursive_pickle(obj, max_depth, path=None, protocol=None, **kwargs): + """Recursively try to pickle an object, replacing unpicklable parts with placeholders. + + Args: + obj: The object to pickle + max_depth: Maximum recursion depth + path: Current path in the object graph + protocol: The pickle protocol version to use + **kwargs: Additional arguments for pickle/dill.dumps + + Returns: + bytes: Pickled data with placeholders for unpicklable objects + """ + if path is None: + path = [] + + obj_type = type(obj) + + # Check if this type is known to be unpicklable + if obj_type in PicklePatcher._unpicklable_types: + placeholder = PicklePatcher._create_placeholder( + obj, + "Known unpicklable type", + path + ) + return dill.dumps(placeholder, protocol=protocol, **kwargs) + + # Check for max depth + if max_depth <= 0: + placeholder = PicklePatcher._create_placeholder( + obj, + "Max recursion depth exceeded", + path + ) + return dill.dumps(placeholder, protocol=protocol, **kwargs) + + # Try standard pickling + success, result = PicklePatcher._pickle(obj, path, protocol, **kwargs) + if success: + return result + + error_msg = result # Error message from pickling attempt + + # Handle different container types + if isinstance(obj, dict): + return PicklePatcher._handle_dict(obj, max_depth, error_msg, path, protocol=protocol, **kwargs) + elif isinstance(obj, (list, tuple, set)): + return PicklePatcher._handle_sequence(obj, max_depth, error_msg, path, protocol=protocol, **kwargs) + elif hasattr(obj, "__dict__"): + result = PicklePatcher._handle_object(obj, max_depth, error_msg, path, protocol=protocol, **kwargs) + + # If this was a failure, add the type to the cache + unpickled = dill.loads(result) + if isinstance(unpickled, PicklePlaceholder): + PicklePatcher._unpicklable_types.add(obj_type) + return result + + # For other unpicklable objects, use a placeholder + placeholder = PicklePatcher._create_placeholder(obj, error_msg, path) + return dill.dumps(placeholder, protocol=protocol, **kwargs) + + @staticmethod + def _handle_dict(obj_dict, max_depth, error_msg, path, protocol=None, **kwargs): + """Handle pickling for dictionary objects. + + Args: + obj_dict: The dictionary to pickle + max_depth: Maximum recursion depth + error_msg: Error message from the original pickling attempt + path: Current path in the object graph + protocol: The pickle protocol version to use + **kwargs: Additional arguments for pickle/dill.dumps + + Returns: + bytes: Pickled data with placeholders for unpicklable objects + """ + if not isinstance(obj_dict, dict): + placeholder = PicklePatcher._create_placeholder( + obj_dict, + f"Expected a dictionary, got {type(obj_dict).__name__}", + path + ) + return dill.dumps(placeholder, protocol=protocol, **kwargs) + + result = {} + + for key, value in obj_dict.items(): + # Process the key + key_success, key_result = PicklePatcher._pickle(key, path, protocol, **kwargs) + if key_success: + key_result = key + else: + # If the key can't be pickled, use a string representation + try: + key_str = str(key)[:50] + except: + key_str = f"" + key_result = f"" + + # Process the value + value_path = path + [f"[{repr(key)[:20]}]"] + value_success, value_bytes = PicklePatcher._pickle(value, value_path, protocol, **kwargs) + + if value_success: + value_result = value + else: + # Try recursive pickling for the value + try: + value_bytes = PicklePatcher._recursive_pickle( + value, max_depth - 1, value_path, protocol=protocol, **kwargs + ) + value_result = dill.loads(value_bytes) + except Exception as inner_e: + value_result = PicklePatcher._create_placeholder( + value, + str(inner_e), + value_path + ) + + result[key_result] = value_result + + return dill.dumps(result, protocol=protocol, **kwargs) + + @staticmethod + def _handle_sequence(obj_seq, max_depth, error_msg, path, protocol=None, **kwargs): + """Handle pickling for sequence types (list, tuple, set). + + Args: + obj_seq: The sequence to pickle + max_depth: Maximum recursion depth + error_msg: Error message from the original pickling attempt + path: Current path in the object graph + protocol: The pickle protocol version to use + **kwargs: Additional arguments for pickle/dill.dumps + + Returns: + bytes: Pickled data with placeholders for unpicklable objects + """ + result = [] + + for i, item in enumerate(obj_seq): + item_path = path + [f"[{i}]"] + + # Try to pickle the item directly + success, _ = PicklePatcher._pickle(item, item_path, protocol, **kwargs) + if success: + result.append(item) + continue + + # If we couldn't pickle directly, try recursively + try: + item_bytes = PicklePatcher._recursive_pickle( + item, max_depth - 1, item_path, protocol=protocol, **kwargs + ) + result.append(dill.loads(item_bytes)) + except Exception as inner_e: + # If recursive pickling fails, use a placeholder + placeholder = PicklePatcher._create_placeholder( + item, + str(inner_e), + item_path + ) + result.append(placeholder) + + # Convert back to the original type + if isinstance(obj_seq, tuple): + result = tuple(result) + elif isinstance(obj_seq, set): + # Try to create a set from the result + try: + result = set(result) + except Exception: + # If we can't create a set (unhashable items), keep it as a list + pass + + return dill.dumps(result, protocol=protocol, **kwargs) + + @staticmethod + def _handle_object(obj, max_depth, error_msg, path, protocol=None, **kwargs): + """Handle pickling for custom objects with __dict__. + + Args: + obj: The object to pickle + max_depth: Maximum recursion depth + error_msg: Error message from the original pickling attempt + path: Current path in the object graph + protocol: The pickle protocol version to use + **kwargs: Additional arguments for pickle/dill.dumps + + Returns: + bytes: Pickled data with placeholders for unpicklable objects + """ + # Try to create a new instance of the same class + try: + # First try to create an empty instance + new_obj = object.__new__(type(obj)) + + # Handle __dict__ attributes if they exist + if hasattr(obj, "__dict__"): + for attr_name, attr_value in obj.__dict__.items(): + attr_path = path + [attr_name] + + # Try to pickle directly first + success, _ = PicklePatcher._pickle(attr_value, attr_path, protocol, **kwargs) + if success: + setattr(new_obj, attr_name, attr_value) + continue + + # If direct pickling fails, try recursive pickling + try: + attr_bytes = PicklePatcher._recursive_pickle( + attr_value, max_depth - 1, attr_path, protocol=protocol, **kwargs + ) + setattr(new_obj, attr_name, dill.loads(attr_bytes)) + except Exception as inner_e: + # Use placeholder for unpicklable attribute + placeholder = PicklePatcher._create_placeholder( + attr_value, + str(inner_e), + attr_path + ) + setattr(new_obj, attr_name, placeholder) + + # Try to pickle the patched object + success, result = PicklePatcher._pickle(new_obj, path, protocol, **kwargs) + if success: + return result + # Fall through to placeholder creation + except Exception: + pass # Fall through to placeholder creation + + # If we get here, just use a placeholder + placeholder = PicklePatcher._create_placeholder(obj, error_msg, path) + return dill.dumps(placeholder, protocol=protocol, **kwargs) \ No newline at end of file diff --git a/codeflash/picklepatch/pickle_placeholder.py b/codeflash/picklepatch/pickle_placeholder.py new file mode 100644 index 000000000..cddb6535a --- /dev/null +++ b/codeflash/picklepatch/pickle_placeholder.py @@ -0,0 +1,66 @@ +class PicklePlaceholder: + """A placeholder for an object that couldn't be pickled. + + When unpickled, any attempt to access attributes or call methods on this + placeholder will raise an informative exception. + """ + + def __init__(self, obj_type, obj_str, error_msg, path=None): + """Initialize a placeholder for an unpicklable object. + + Args: + obj_type (str): The type name of the original object + obj_str (str): String representation of the original object + error_msg (str): The error message that occurred during pickling + path (list, optional): Path to this object in the original object graph + + """ + # Store these directly in __dict__ to avoid __getattr__ recursion + self.__dict__["obj_type"] = obj_type + self.__dict__["obj_str"] = obj_str + self.__dict__["error_msg"] = error_msg + self.__dict__["path"] = path if path is not None else [] + + def __getattr__(self, name): + """Raise an error when any attribute is accessed.""" + path_str = ".".join(self.__dict__["path"]) if self.__dict__["path"] else "root object" + raise AttributeError( + f"Cannot access attribute '{name}' on unpicklable object at {path_str}. " + f"Original type: {self.__dict__['obj_type']}. Error: {self.__dict__['error_msg']}" + ) + + def __setattr__(self, name, value): + """Prevent setting attributes.""" + self.__getattr__(name) # This will raise an AttributeError + + def __call__(self, *args, **kwargs): + """Raise an error when the object is called.""" + path_str = ".".join(self.__dict__["path"]) if self.__dict__["path"] else "root object" + raise TypeError( + f"Cannot call unpicklable object at {path_str}. " + f"Original type: {self.__dict__['obj_type']}. Error: {self.__dict__['error_msg']}" + ) + + def __repr__(self): + """Return a string representation of the placeholder.""" + try: + path_str = ".".join(self.__dict__["path"]) if self.__dict__["path"] else "root" + return f"" + except: + return "" + + def __str__(self): + """Return a string representation of the placeholder.""" + return self.__repr__() + + def __reduce__(self): + """Make sure pickling of the placeholder itself works correctly.""" + return ( + PicklePlaceholder, + ( + self.__dict__["obj_type"], + self.__dict__["obj_str"], + self.__dict__["error_msg"], + self.__dict__["path"] + ) + ) diff --git a/tests/test_pickle_patcher.py b/tests/test_pickle_patcher.py new file mode 100644 index 000000000..a85dce047 --- /dev/null +++ b/tests/test_pickle_patcher.py @@ -0,0 +1,172 @@ + +import socket +import pytest +import requests +import sqlite3 + +try: + import sqlalchemy + from sqlalchemy.orm import Session + from sqlalchemy import create_engine, Column, Integer, String + from sqlalchemy.ext.declarative import declarative_base + + HAS_SQLALCHEMY = True +except ImportError: + HAS_SQLALCHEMY = False + +from codeflash.picklepatch.pickle_patcher import PicklePatcher +from codeflash.picklepatch.pickle_placeholder import PicklePlaceholder +def test_picklepatch_simple_nested(): + """ + Test that a simple nested data structure pickles and unpickles correctly. + """ + original_data = { + "numbers": [1, 2, 3], + "nested_dict": {"key": "value", "another": 42}, + } + + dumped = PicklePatcher.dumps(original_data) + reloaded = PicklePatcher.loads(dumped) + + assert reloaded == original_data + # Everything was pickleable, so no placeholders should appear. + +def test_picklepatch_with_socket(): + """ + Test that a data structure containing a raw socket is replaced by + PicklePlaceholder rather than raising an error. + """ + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + data_with_socket = { + "safe_value": 123, + "raw_socket": s, + } + + dumped = PicklePatcher.dumps(data_with_socket) + reloaded = PicklePatcher.loads(dumped) + + # We expect "raw_socket" to be replaced by a placeholder + assert isinstance(reloaded, dict) + assert reloaded["safe_value"] == 123 + assert isinstance(reloaded["raw_socket"], PicklePlaceholder) + + # Attempting to use or access attributes => AttributeError + # (not RuntimeError as in original tests, our implementation uses AttributeError) + with pytest.raises(AttributeError): + reloaded["raw_socket"].recv(1024) + + +def test_picklepatch_deeply_nested(): + """ + Test that deep nesting with unpicklable objects works correctly. + """ + # Create a deeply nested structure with an unpicklable object + deep_nested = { + "level1": { + "level2": { + "level3": { + "normal": "value", + "socket": socket.socket(socket.AF_INET, socket.SOCK_STREAM) + } + } + } + } + + dumped = PicklePatcher.dumps(deep_nested) + reloaded = PicklePatcher.loads(dumped) + + # We should be able to access the normal value + assert reloaded["level1"]["level2"]["level3"]["normal"] == "value" + + # The socket should be replaced with a placeholder + assert isinstance(reloaded["level1"]["level2"]["level3"]["socket"], PicklePlaceholder) + +def test_picklepatch_class_with_unpicklable_attr(): + """ + Test that a class with an unpicklable attribute works correctly. + """ + class TestClass: + def __init__(self): + self.normal = "normal value" + self.unpicklable = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + + obj = TestClass() + + dumped = PicklePatcher.dumps(obj) + reloaded = PicklePatcher.loads(dumped) + + # Normal attribute should be preserved + assert reloaded.normal == "normal value" + + # Unpicklable attribute should be replaced with a placeholder + assert isinstance(reloaded.unpicklable, PicklePlaceholder) + + + + +def test_picklepatch_with_database_connection(): + """ + Test that a data structure containing a database connection is replaced + by PicklePlaceholder rather than raising an error. + """ + # SQLite connection - not pickleable + conn = sqlite3.connect(':memory:') + cursor = conn.cursor() + + data_with_db = { + "description": "Database connection", + "connection": conn, + "cursor": cursor, + } + + dumped = PicklePatcher.dumps(data_with_db) + reloaded = PicklePatcher.loads(dumped) + + # Both connection and cursor should become placeholders + assert isinstance(reloaded, dict) + assert reloaded["description"] == "Database connection" + assert isinstance(reloaded["connection"], PicklePlaceholder) + assert isinstance(reloaded["cursor"], PicklePlaceholder) + + # Attempting to use attributes => AttributeError + with pytest.raises(AttributeError): + reloaded["connection"].execute("SELECT 1") + + +def test_picklepatch_with_generator(): + """ + Test that a data structure containing a generator is replaced by + PicklePlaceholder rather than raising an error. + """ + + def simple_generator(): + yield 1 + yield 2 + yield 3 + + # Create a generator + gen = simple_generator() + + # Put it in a data structure + data_with_generator = { + "description": "Contains a generator", + "generator": gen, + "normal_list": [1, 2, 3] + } + + dumped = PicklePatcher.dumps(data_with_generator) + reloaded = PicklePatcher.loads(dumped) + + # Generator should be replaced with a placeholder + assert isinstance(reloaded, dict) + assert reloaded["description"] == "Contains a generator" + assert reloaded["normal_list"] == [1, 2, 3] + assert isinstance(reloaded["generator"], PicklePlaceholder) + + # Attempting to use the generator => AttributeError + with pytest.raises(TypeError): + next(reloaded["generator"]) + + # Attempting to call methods on the generator => AttributeError + with pytest.raises(AttributeError): + reloaded["generator"].send(None) \ No newline at end of file From d653d0dc4290dabbfb8107b5791762263d5c8f86 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 8 Apr 2025 13:32:21 -0700 Subject: [PATCH 053/122] draft of end to end test --- code_to_optimize/bubble_sort_picklepatch.py | 78 +++++++++ .../pytest/test_bubble_sort_picklepatch.py | 34 ++++ codeflash/verification/parse_test_output.py | 6 +- tests/test_pickle_patcher.py | 153 +++++++++++++++++- 4 files changed, 265 insertions(+), 6 deletions(-) create mode 100644 code_to_optimize/bubble_sort_picklepatch.py create mode 100644 code_to_optimize/tests/pytest/test_bubble_sort_picklepatch.py diff --git a/code_to_optimize/bubble_sort_picklepatch.py b/code_to_optimize/bubble_sort_picklepatch.py new file mode 100644 index 000000000..25cbe9628 --- /dev/null +++ b/code_to_optimize/bubble_sort_picklepatch.py @@ -0,0 +1,78 @@ +def bubble_sort_with_unused_socket(data_container): + """ + Performs a bubble sort on a list within the data_container. The data container has the following schema: + - 'numbers' (list): The list to be sorted. + - 'socket' (socket): A socket + + Args: + data_container: A dictionary with at least 'numbers' (list) and 'socket' keys + + Returns: + list: The sorted list of numbers + """ + # Extract the list to sort, leaving the socket untouched + numbers = data_container.get('numbers', []).copy() + + # Classic bubble sort implementation + n = len(numbers) + for i in range(n): + # Flag to optimize by detecting if no swaps occurred + swapped = False + + # Last i elements are already in place + for j in range(0, n - i - 1): + # Swap if the element is greater than the next element + if numbers[j] > numbers[j + 1]: + numbers[j], numbers[j + 1] = numbers[j + 1], numbers[j] + swapped = True + + # If no swapping occurred in this pass, the list is sorted + if not swapped: + break + + return numbers + + +def bubble_sort_with_used_socket(data_container): + """ + Performs a bubble sort on a list within the data_container. The data container has the following schema: + - 'numbers' (list): The list to be sorted. + - 'socket' (socket): A socket + + Args: + data_container: A dictionary with at least 'numbers' (list) and 'socket' keys + + Returns: + list: The sorted list of numbers + """ + # Extract the list to sort and socket + numbers = data_container.get('numbers', []).copy() + socket = data_container.get('socket') + + # Track swap count + swap_count = 0 + + # Classic bubble sort implementation + n = len(numbers) + for i in range(n): + # Flag to optimize by detecting if no swaps occurred + swapped = False + + # Last i elements are already in place + for j in range(0, n - i - 1): + # Swap if the element is greater than the next element + if numbers[j] > numbers[j + 1]: + # Perform the swap + numbers[j], numbers[j + 1] = numbers[j + 1], numbers[j] + swapped = True + swap_count += 1 + + # If no swapping occurred in this pass, the list is sorted + if not swapped: + break + + # Send final summary + summary = f"Bubble sort completed with {swap_count} swaps" + socket.send(summary.encode()) + + return numbers \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/test_bubble_sort_picklepatch.py b/code_to_optimize/tests/pytest/test_bubble_sort_picklepatch.py new file mode 100644 index 000000000..9f3e0f9af --- /dev/null +++ b/code_to_optimize/tests/pytest/test_bubble_sort_picklepatch.py @@ -0,0 +1,34 @@ +import socket +from unittest.mock import Mock + +import pytest + +from code_to_optimize.bubble_sort_picklepatch import bubble_sort_with_unused_socket, bubble_sort_with_used_socket + + +def test_bubble_sort_with_unused_socket(): + mock_socket = Mock() + # Test case 1: Regular unsorted list + data_container = { + 'numbers': [5, 2, 9, 1, 5, 6], + 'socket': mock_socket + } + + result = bubble_sort_with_unused_socket(data_container) + + # Check that the result is correctly sorted + assert result == [1, 2, 5, 5, 6, 9] + +def test_bubble_sort_with_used_socket(): + mock_socket = Mock() + # Test case 1: Regular unsorted list + data_container = { + 'numbers': [5, 2, 9, 1, 5, 6], + 'socket': mock_socket + } + + result = bubble_sort_with_used_socket(data_container) + + # Check that the result is correctly sorted + assert result == [1, 2, 5, 5, 6, 9] + diff --git a/codeflash/verification/parse_test_output.py b/codeflash/verification/parse_test_output.py index 924e2876a..cc29b6cda 100644 --- a/codeflash/verification/parse_test_output.py +++ b/codeflash/verification/parse_test_output.py @@ -8,7 +8,6 @@ from pathlib import Path from typing import TYPE_CHECKING -import dill as pickle from junitparser.xunit2 import JUnitXml from lxml.etree import XMLParser, parse @@ -21,6 +20,7 @@ ) from codeflash.discovery.discover_unit_tests import discover_parameters_unittest from codeflash.models.models import FunctionTestInvocation, InvocationId, TestResults, TestType, VerificationType +from codeflash.picklepatch.pickle_patcher import PicklePatcher from codeflash.verification.coverage_utils import CoverageUtils if TYPE_CHECKING: @@ -75,7 +75,7 @@ def parse_test_return_values_bin(file_location: Path, test_files: TestFiles, tes test_type = test_files.get_test_type_by_instrumented_file_path(test_file_path) try: - test_pickle = pickle.loads(test_pickle_bin) if loop_index == 1 else None + test_pickle = PicklePatcher.loads(test_pickle_bin) if loop_index == 1 else None except Exception as e: if DEBUG_MODE: logger.exception(f"Failed to load pickle file for {encoded_test_name} Exception: {e}") @@ -133,7 +133,7 @@ def parse_sqlite_test_results(sqlite_file_path: Path, test_files: TestFiles, tes # TODO : this is because sqlite writes original file module path. Should make it consistent test_type = test_files.get_test_type_by_original_file_path(test_file_path) try: - ret_val = (pickle.loads(val[7]) if loop_index == 1 else None,) + ret_val = (PicklePatcher.loads(val[7]) if loop_index == 1 else None,) except Exception: continue test_results.add( diff --git a/tests/test_pickle_patcher.py b/tests/test_pickle_patcher.py index a85dce047..05bd06f15 100644 --- a/tests/test_pickle_patcher.py +++ b/tests/test_pickle_patcher.py @@ -1,9 +1,19 @@ - +import os +import pickle import socket +from argparse import Namespace +from pathlib import Path + +import dill import pytest import requests import sqlite3 +from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test +from codeflash.discovery.functions_to_optimize import FunctionToOptimize +from codeflash.models.models import CodePosition, TestingMode, TestType, TestFiles, TestFile +from codeflash.optimization.optimizer import Optimizer + try: import sqlalchemy from sqlalchemy.orm import Session @@ -52,7 +62,7 @@ def test_picklepatch_with_socket(): # Attempting to use or access attributes => AttributeError # (not RuntimeError as in original tests, our implementation uses AttributeError) - with pytest.raises(AttributeError): + with pytest.raises(AttributeError) : reloaded["raw_socket"].recv(1024) @@ -169,4 +179,141 @@ def simple_generator(): # Attempting to call methods on the generator => AttributeError with pytest.raises(AttributeError): - reloaded["generator"].send(None) \ No newline at end of file + reloaded["generator"].send(None) + + +def test_picklepatch_loads_standard_pickle(): + """ + Test that PicklePatcher.loads can correctly load data that was pickled + using the standard pickle module. + """ + # Create a simple data structure + original_data = { + "numbers": [1, 2, 3], + "nested_dict": {"key": "value", "another": 42}, + "tuple": (1, "two", 3.0), + } + + # Pickle it with standard pickle + pickled_data = pickle.dumps(original_data) + + # Load with PicklePatcher + reloaded = PicklePatcher.loads(pickled_data) + + # Verify the data is correctly loaded + assert reloaded == original_data + assert isinstance(reloaded, dict) + assert reloaded["numbers"] == [1, 2, 3] + assert reloaded["nested_dict"]["key"] == "value" + assert reloaded["tuple"] == (1, "two", 3.0) + + +def test_picklepatch_loads_dill_pickle(): + """ + Test that PicklePatcher.loads can correctly load data that was pickled + using the dill module, which can pickle more complex objects than the + standard pickle module. + """ + + # Create a more complex data structure that includes a lambda function + # which dill can handle but standard pickle cannot + original_data = { + "numbers": [1, 2, 3], + "function": lambda x: x * 2, + "nested": { + "another_function": lambda y: y ** 2 + } + } + + # Pickle it with dill + dilled_data = dill.dumps(original_data) + + # Load with PicklePatcher + reloaded = PicklePatcher.loads(dilled_data) + + # Verify the data structure + assert isinstance(reloaded, dict) + assert reloaded["numbers"] == [1, 2, 3] + + # Test that the functions actually work + assert reloaded["function"](5) == 10 + assert reloaded["nested"]["another_function"](4) == 16 + +def test_run_and_parse_picklepatch() -> None: + + test_path = ( + Path(__file__).parent.resolve() + / "../code_to_optimize/tests/pytest/test_bubble_sort_picklepatch.py" + ).resolve() + test_path_perf = ( + Path(__file__).parent.resolve() + / "../code_to_optimize/tests/pytest/test_bubble_sort_picklepatch_perf.py" + ).resolve() + fto_path = (Path(__file__).parent.resolve() / "../code_to_optimize/bubble_sort_picklepatch.py").resolve() + original_test =test_path.read_text("utf-8") + try: + tests_root = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/").resolve() + project_root_path = (Path(__file__).parent / "..").resolve() + original_cwd = Path.cwd() + run_cwd = Path(__file__).parent.parent.resolve() + func = FunctionToOptimize(function_name="bubble_sort_with_unused_socket", parents=[], file_path=Path(fto_path)) + os.chdir(run_cwd) + success, new_test = inject_profiling_into_existing_test( + test_path, + [CodePosition(13,14), CodePosition(31,14)], + func, + project_root_path, + "pytest", + mode=TestingMode.BEHAVIOR, + ) + os.chdir(original_cwd) + assert success + assert new_test is not None + + with test_path.open("w") as f: + f.write(new_test) + + opt = Optimizer( + Namespace( + project_root=project_root_path, + disable_telemetry=True, + tests_root=tests_root, + test_framework="pytest", + pytest_cmd="pytest", + experiment_id=None, + test_project_root=project_root_path, + ) + ) + test_env = os.environ.copy() + test_env["CODEFLASH_TEST_ITERATION"] = "0" + test_env["CODEFLASH_LOOP_INDEX"] = "1" + test_type = TestType.EXISTING_UNIT_TEST + + func_optimizer = opt.create_function_optimizer(func) + func_optimizer.test_files = TestFiles( + test_files=[ + TestFile( + instrumented_behavior_file_path=test_path, + test_type=test_type, + original_file_path=test_path, + benchmarking_file_path=test_path_perf, + ) + ] + ) + test_results, coverage_data = func_optimizer.run_and_parse_tests( + testing_type=TestingMode.BEHAVIOR, + test_env=test_env, + test_files=func_optimizer.test_files, + optimization_iteration=0, + pytest_min_loops=1, + pytest_max_loops=1, + testing_time=0.1, + ) + assert test_results.test_results[0].id.test_function_name =="test_bubble_sort_with_unused_socket" + assert test_results.test_results[0].did_pass ==True + assert test_results.test_results[1].id.test_function_name =="test_bubble_sort_with_used_socket" + assert test_results.test_results[1].did_pass ==False + # assert pickle placeholder problem + print(test_results) + finally: + test_path.write_text(original_test) \ No newline at end of file From a73b541159b60d07b64cbfb60c99d3501375881f Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 25 Feb 2025 13:18:08 -0800 Subject: [PATCH 054/122] initial implementation for pytest benchmark discovery --- .../pytest/test_benchmark_bubble_sort.py | 6 + codeflash/discovery/discover_unit_tests.py | 2 + .../discovery/pytest_new_process_discovery.py | 12 +- codeflash/verification/test_results.py | 242 ++++++++++++++++++ tests/test_unit_test_discovery.py | 14 + 5 files changed, 275 insertions(+), 1 deletion(-) create mode 100644 code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py create mode 100644 codeflash/verification/test_results.py diff --git a/code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py new file mode 100644 index 000000000..dcbb86ac1 --- /dev/null +++ b/code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py @@ -0,0 +1,6 @@ +from code_to_optimize.bubble_sort import sorter + + +def test_sort(benchmark): + result = benchmark(sorter, list(reversed(range(5000)))) + assert result == list(range(5000)) diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py index e26680e1a..4820d6e65 100644 --- a/codeflash/discovery/discover_unit_tests.py +++ b/codeflash/discovery/discover_unit_tests.py @@ -107,6 +107,8 @@ def discover_tests_pytest( test_type = TestType.REPLAY_TEST elif "test_concolic_coverage" in test["test_file"]: test_type = TestType.CONCOLIC_COVERAGE_TEST + elif test["test_type"] == "benchmark": # New condition for benchmark tests + test_type = TestType.BENCHMARK_TEST else: test_type = TestType.EXISTING_UNIT_TEST diff --git a/codeflash/discovery/pytest_new_process_discovery.py b/codeflash/discovery/pytest_new_process_discovery.py index 2d8583255..fa2dea5e9 100644 --- a/codeflash/discovery/pytest_new_process_discovery.py +++ b/codeflash/discovery/pytest_new_process_discovery.py @@ -29,7 +29,17 @@ def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, s test_class = None if test.cls: test_class = test.parent.name - test_results.append({"test_file": str(test.path), "test_class": test_class, "test_function": test.name}) + + # Determine if this is a benchmark test by checking for the benchmark fixture + is_benchmark = hasattr(test, 'fixturenames') and 'benchmark' in test.fixturenames + test_type = 'benchmark' if is_benchmark else 'regular' + + test_results.append({ + "test_file": str(test.path), + "test_class": test_class, + "test_function": test.name, + "test_type": test_type + }) return test_results diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py new file mode 100644 index 000000000..99151f983 --- /dev/null +++ b/codeflash/verification/test_results.py @@ -0,0 +1,242 @@ +from __future__ import annotations + +import sys +from collections.abc import Iterator +from enum import Enum +from pathlib import Path +from typing import Optional, cast + +from pydantic import BaseModel +from pydantic.dataclasses import dataclass +from rich.tree import Tree + +from codeflash.cli_cmds.console import DEBUG_MODE, logger +from codeflash.verification.comparator import comparator + + +class VerificationType(str, Enum): + FUNCTION_CALL = ( + "function_call" # Correctness verification for a test function, checks input values and output values) + ) + INIT_STATE_FTO = "init_state_fto" # Correctness verification for fto class instance attributes after init + INIT_STATE_HELPER = "init_state_helper" # Correctness verification for helper class instance attributes after init + + +class TestType(Enum): + EXISTING_UNIT_TEST = 1 + INSPIRED_REGRESSION = 2 + GENERATED_REGRESSION = 3 + REPLAY_TEST = 4 + CONCOLIC_COVERAGE_TEST = 5 + INIT_STATE_TEST = 6 + BENCHMARK_TEST = 7 + + def to_name(self) -> str: + if self == TestType.INIT_STATE_TEST: + return "" + names = { + TestType.EXISTING_UNIT_TEST: "βš™οΈ Existing Unit Tests", + TestType.INSPIRED_REGRESSION: "🎨 Inspired Regression Tests", + TestType.GENERATED_REGRESSION: "πŸŒ€ Generated Regression Tests", + TestType.REPLAY_TEST: "βͺ Replay Tests", + TestType.CONCOLIC_COVERAGE_TEST: "πŸ”Ž Concolic Coverage Tests", + TestType.BENCHMARK_TEST: "πŸ“ Benchmark Tests", + } + return names[self] + + +@dataclass(frozen=True) +class InvocationId: + test_module_path: str # The fully qualified name of the test module + test_class_name: Optional[str] # The name of the class where the test is defined + test_function_name: Optional[str] # The name of the test_function. Does not include the components of the file_name + function_getting_tested: str + iteration_id: Optional[str] + + # test_module_path:TestSuiteClass.test_function_name:function_tested:iteration_id + def id(self) -> str: + return f"{self.test_module_path}:{(self.test_class_name + '.' if self.test_class_name else '')}{self.test_function_name}:{self.function_getting_tested}:{self.iteration_id}" + + @staticmethod + def from_str_id(string_id: str, iteration_id: Optional[str] = None) -> InvocationId: + components = string_id.split(":") + assert len(components) == 4 + second_components = components[1].split(".") + if len(second_components) == 1: + test_class_name = None + test_function_name = second_components[0] + else: + test_class_name = second_components[0] + test_function_name = second_components[1] + # logger.debug(f"Invocation id info: test_module_path: {components[0]}, test_class_name: {test_class_name}, test_function_name: {test_function_name}, function_getting_tested: {components[2]}, iteration_id: {iteration_id if iteration_id else components[3]}") + return InvocationId( + test_module_path=components[0], + test_class_name=test_class_name, + test_function_name=test_function_name, + function_getting_tested=components[2], + iteration_id=iteration_id if iteration_id else components[3], + ) + + +@dataclass(frozen=True) +class FunctionTestInvocation: + loop_index: int # The loop index of the function invocation, starts at 1 + id: InvocationId # The fully qualified name of the function invocation (id) + file_name: Path # The file where the test is defined + did_pass: bool # Whether the test this function invocation was part of, passed or failed + runtime: Optional[int] # Time in nanoseconds + test_framework: str # unittest or pytest + test_type: TestType + return_value: Optional[object] # The return value of the function invocation + timed_out: Optional[bool] + verification_type: Optional[str] = VerificationType.FUNCTION_CALL + + @property + def unique_invocation_loop_id(self) -> str: + return f"{self.loop_index}:{self.id.id()}" + + +class TestResults(BaseModel): + # don't modify these directly, use the add method + # also we don't support deletion of test results elements - caution is advised + test_results: list[FunctionTestInvocation] = [] + test_result_idx: dict[str, int] = {} + + def add(self, function_test_invocation: FunctionTestInvocation) -> None: + unique_id = function_test_invocation.unique_invocation_loop_id + if unique_id in self.test_result_idx: + if DEBUG_MODE: + logger.warning(f"Test result with id {unique_id} already exists. SKIPPING") + return + self.test_result_idx[unique_id] = len(self.test_results) + self.test_results.append(function_test_invocation) + + def merge(self, other: TestResults) -> None: + original_len = len(self.test_results) + self.test_results.extend(other.test_results) + for k, v in other.test_result_idx.items(): + if k in self.test_result_idx: + msg = f"Test result with id {k} already exists." + raise ValueError(msg) + self.test_result_idx[k] = v + original_len + + def get_by_unique_invocation_loop_id(self, unique_invocation_loop_id: str) -> FunctionTestInvocation | None: + try: + return self.test_results[self.test_result_idx[unique_invocation_loop_id]] + except (IndexError, KeyError): + return None + + def get_all_ids(self) -> set[InvocationId]: + return {test_result.id for test_result in self.test_results} + + def get_all_unique_invocation_loop_ids(self) -> set[str]: + return {test_result.unique_invocation_loop_id for test_result in self.test_results} + + def number_of_loops(self) -> int: + if not self.test_results: + return 0 + return max(test_result.loop_index for test_result in self.test_results) + + def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]: + report = {} + for test_type in TestType: + report[test_type] = {"passed": 0, "failed": 0} + for test_result in self.test_results: + if test_result.loop_index == 1: + if test_result.did_pass: + report[test_result.test_type]["passed"] += 1 + else: + report[test_result.test_type]["failed"] += 1 + return report + + @staticmethod + def report_to_string(report: dict[TestType, dict[str, int]]) -> str: + return " ".join( + [ + f"{test_type.to_name()}- (Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']})" + for test_type in TestType + ] + ) + + @staticmethod + def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree: + tree = Tree(title) + for test_type in TestType: + tree.add( + f"{test_type.to_name()} - Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']}" + ) + return tree + + def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: + for result in self.test_results: + if result.did_pass and not result.runtime: + pass + # logger.debug( + # f"Ignoring test case that passed but had no runtime -> {result.id}, Loop # {result.loop_index}, Test Type: {result.test_type}, Verification Type: {result.verification_type}" + # ) + usable_runtimes = [ + (result.id, result.runtime) for result in self.test_results if result.did_pass and result.runtime + ] + return { + usable_id: [runtime[1] for runtime in usable_runtimes if runtime[0] == usable_id] + for usable_id in {runtime[0] for runtime in usable_runtimes} + } + + def total_passed_runtime(self) -> int: + """Calculate the sum of runtimes of all test cases that passed, where a testcase runtime + is the minimum value of all looped execution runtimes. + + :return: The runtime in nanoseconds. + """ + return sum( + [ + min(usable_runtime_data) + for invocation_id, usable_runtime_data in self.usable_runtime_data_by_test_case().items() + ] + ) + + def __iter__(self) -> Iterator[FunctionTestInvocation]: + return iter(self.test_results) + + def __len__(self) -> int: + return len(self.test_results) + + def __getitem__(self, index: int) -> FunctionTestInvocation: + return self.test_results[index] + + def __setitem__(self, index: int, value: FunctionTestInvocation) -> None: + self.test_results[index] = value + + def __contains__(self, value: FunctionTestInvocation) -> bool: + return value in self.test_results + + def __bool__(self) -> bool: + return bool(self.test_results) + + def __eq__(self, other: object) -> bool: + # Unordered comparison + if type(self) is not type(other): + return False + if len(self) != len(other): + return False + original_recursion_limit = sys.getrecursionlimit() + cast(TestResults, other) + for test_result in self: + other_test_result = other.get_by_unique_invocation_loop_id(test_result.unique_invocation_loop_id) + if other_test_result is None: + return False + + if original_recursion_limit < 5000: + sys.setrecursionlimit(5000) + if ( + test_result.file_name != other_test_result.file_name + or test_result.did_pass != other_test_result.did_pass + or test_result.runtime != other_test_result.runtime + or test_result.test_framework != other_test_result.test_framework + or test_result.test_type != other_test_result.test_type + or not comparator(test_result.return_value, other_test_result.return_value) + ): + sys.setrecursionlimit(original_recursion_limit) + return False + sys.setrecursionlimit(original_recursion_limit) + return True diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index 73556928e..1aad04d42 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -19,6 +19,20 @@ def test_unit_test_discovery_pytest(): assert len(tests) > 0 # print(tests) +def test_benchmark_test_discovery_pytest(): + project_path = Path(__file__).parent.parent.resolve() / "code_to_optimize" + tests_path = project_path / "tests" / "pytest" + test_config = TestConfig( + tests_root=tests_path, + project_root_path=project_path, + test_framework="pytest", + tests_project_rootdir=tests_path.parent, + ) + tests = discover_unit_tests(test_config) + print(tests) + assert len(tests) > 0 + # print(tests) + def test_unit_test_discovery_unittest(): project_path = Path(__file__).parent.parent.resolve() / "code_to_optimize" From 965e2c818cacb0f73f5361a462ca73aef58465d9 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 27 Feb 2025 16:25:07 -0800 Subject: [PATCH 055/122] initial implementation for tracing benchmarks using a plugin, and projecting speedup --- code_to_optimize/process_and_bubble_sort.py | 28 + .../test_benchmark_bubble_sort.py | 7 + .../benchmarks/test_process_and_sort.py | 8 + codeflash/benchmarking/__init__.py | 0 codeflash/benchmarking/get_trace_info.py | 112 ++++ codeflash/benchmarking/plugin/__init__.py | 0 codeflash/benchmarking/plugin/plugin.py | 79 +++ .../pytest_new_process_trace_benchmarks.py | 15 + codeflash/benchmarking/trace_benchmarks.py | 20 + codeflash/cli_cmds/cli.py | 33 +- codeflash/discovery/discover_unit_tests.py | 2 - codeflash/discovery/functions_to_optimize.py | 6 +- .../pytest_new_process_discover_benchmarks.py | 54 ++ .../discovery/pytest_new_process_discovery.py | 12 +- codeflash/optimization/function_optimizer.py | 380 ++++++------- codeflash/optimization/optimizer.py | 51 +- codeflash/tracer.py | 507 +++++++----------- codeflash/verification/test_results.py | 2 - codeflash/verification/verification_utils.py | 1 + pyproject.toml | 30 +- tests/test_trace_benchmarks.py | 8 + tests/test_unit_test_discovery.py | 8 +- 22 files changed, 791 insertions(+), 572 deletions(-) create mode 100644 code_to_optimize/process_and_bubble_sort.py rename code_to_optimize/tests/pytest/{ => benchmarks}/test_benchmark_bubble_sort.py (50%) create mode 100644 code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py create mode 100644 codeflash/benchmarking/__init__.py create mode 100644 codeflash/benchmarking/get_trace_info.py create mode 100644 codeflash/benchmarking/plugin/__init__.py create mode 100644 codeflash/benchmarking/plugin/plugin.py create mode 100644 codeflash/benchmarking/pytest_new_process_trace_benchmarks.py create mode 100644 codeflash/benchmarking/trace_benchmarks.py create mode 100644 codeflash/discovery/pytest_new_process_discover_benchmarks.py create mode 100644 tests/test_trace_benchmarks.py diff --git a/code_to_optimize/process_and_bubble_sort.py b/code_to_optimize/process_and_bubble_sort.py new file mode 100644 index 000000000..94359e599 --- /dev/null +++ b/code_to_optimize/process_and_bubble_sort.py @@ -0,0 +1,28 @@ +from code_to_optimize.bubble_sort import sorter + + +def calculate_pairwise_products(arr): + """ + Calculate the average of all pairwise products in the array. + """ + sum_of_products = 0 + count = 0 + + for i in range(len(arr)): + for j in range(len(arr)): + if i != j: + sum_of_products += arr[i] * arr[j] + count += 1 + + # The average of all pairwise products + return sum_of_products / count if count > 0 else 0 + + +def compute_and_sort(arr): + # Compute pairwise sums average + pairwise_average = calculate_pairwise_products(arr) + + # Call sorter function + sorter(arr.copy()) + + return pairwise_average diff --git a/code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py similarity index 50% rename from code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py rename to code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py index dcbb86ac1..f1ebcf5c7 100644 --- a/code_to_optimize/tests/pytest/test_benchmark_bubble_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py @@ -1,6 +1,13 @@ +import pytest + from code_to_optimize.bubble_sort import sorter def test_sort(benchmark): result = benchmark(sorter, list(reversed(range(5000)))) assert result == list(range(5000)) + +# This should not be picked up as a benchmark test +def test_sort2(): + result = sorter(list(reversed(range(5000)))) + assert result == list(range(5000)) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py new file mode 100644 index 000000000..ca2f0ef65 --- /dev/null +++ b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py @@ -0,0 +1,8 @@ +from code_to_optimize.process_and_bubble_sort import compute_and_sort +from code_to_optimize.bubble_sort2 import sorter +def test_compute_and_sort(benchmark): + result = benchmark(compute_and_sort, list(reversed(range(5000)))) + assert result == 6247083.5 + +def test_no_func(benchmark): + benchmark(sorter, list(reversed(range(5000)))) \ No newline at end of file diff --git a/codeflash/benchmarking/__init__.py b/codeflash/benchmarking/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/codeflash/benchmarking/get_trace_info.py b/codeflash/benchmarking/get_trace_info.py new file mode 100644 index 000000000..1d0b339d9 --- /dev/null +++ b/codeflash/benchmarking/get_trace_info.py @@ -0,0 +1,112 @@ +import sqlite3 +from pathlib import Path +from typing import Dict, Set + +from codeflash.discovery.functions_to_optimize import FunctionToOptimize + + +def get_function_benchmark_timings(trace_dir: Path, all_functions_to_optimize: list[FunctionToOptimize]) -> dict[str, dict[str, float]]: + """Process all trace files in the given directory and extract timing data for the specified functions. + + Args: + trace_dir: Path to the directory containing .trace files + all_functions_to_optimize: Set of FunctionToOptimize objects representing functions to include + + Returns: + A nested dictionary where: + - Outer keys are function qualified names with file name + - Inner keys are benchmark names (trace filename without .trace extension) + - Values are function timing in milliseconds + + """ + # Create a mapping of (filename, function_name, class_name) -> qualified_name for efficient lookups + function_lookup = {} + function_benchmark_timings = {} + + for func in all_functions_to_optimize: + qualified_name = func.qualified_name_with_file_name + + # Extract components (assumes Path.name gives only filename without directory) + filename = func.file_path + function_name = func.function_name + + # Get class name if there's a parent + class_name = func.parents[0].name if func.parents else None + + # Store in lookup dictionary + key = (filename, function_name, class_name) + function_lookup[key] = qualified_name + function_benchmark_timings[qualified_name] = {} + + # Find all .trace files in the directory + trace_files = list(trace_dir.glob("*.trace")) + + for trace_file in trace_files: + # Extract benchmark name from filename (without .trace) + benchmark_name = trace_file.stem + + # Connect to the trace database + conn = sqlite3.connect(trace_file) + cursor = conn.cursor() + + # For each function we're interested in, query the database directly + for (filename, function_name, class_name), qualified_name in function_lookup.items(): + # Adjust query based on whether we have a class name + if class_name: + cursor.execute( + "SELECT total_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND class_name = ?", + (f"%{filename}", function_name, class_name) + ) + else: + cursor.execute( + "SELECT total_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND (class_name IS NULL OR class_name = '')", + (f"%{filename}", function_name) + ) + + result = cursor.fetchone() + if result: + time_ns = result[0] + function_benchmark_timings[qualified_name][benchmark_name] = time_ns / 1e6 # Convert to milliseconds + + conn.close() + + return function_benchmark_timings + + +def get_benchmark_timings(trace_dir: Path) -> dict[str, float]: + """Extract total benchmark timings from trace files. + + Args: + trace_dir: Path to the directory containing .trace files + + Returns: + A dictionary mapping benchmark names to their total execution time in milliseconds. + """ + benchmark_timings = {} + + # Find all .trace files in the directory + trace_files = list(trace_dir.glob("*.trace")) + + for trace_file in trace_files: + # Extract benchmark name from filename (without .trace extension) + benchmark_name = trace_file.stem + + # Connect to the trace database + conn = sqlite3.connect(trace_file) + cursor = conn.cursor() + + # Query the total_time table for the benchmark's total execution time + try: + cursor.execute("SELECT time_ns FROM total_time") + result = cursor.fetchone() + if result: + time_ns = result[0] + # Convert nanoseconds to milliseconds + benchmark_timings[benchmark_name] = time_ns / 1e6 + except sqlite3.OperationalError: + # Handle case where total_time table might not exist + print(f"Warning: Could not get total time for benchmark {benchmark_name}") + + conn.close() + + return benchmark_timings diff --git a/codeflash/benchmarking/plugin/__init__.py b/codeflash/benchmarking/plugin/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py new file mode 100644 index 000000000..34ca2b777 --- /dev/null +++ b/codeflash/benchmarking/plugin/plugin.py @@ -0,0 +1,79 @@ +import pytest + +from codeflash.tracer import Tracer +from pathlib import Path + +class CodeFlashPlugin: + @staticmethod + def pytest_addoption(parser): + parser.addoption( + "--codeflash-trace", + action="store_true", + 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): + if hasattr(plugin, "name") and plugin.name == "pytest-benchmark": + manager.unregister(plugin) + + @staticmethod + def pytest_collection_modifyitems(config, items): + 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) + + @staticmethod + @pytest.fixture + def benchmark(request): + if not request.config.getoption("--codeflash-trace"): + return None + + 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 + ) + + with tracer: + result = func(*args, **kwargs) + + return result + + return Benchmark() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py new file mode 100644 index 000000000..b892d62a0 --- /dev/null +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -0,0 +1,15 @@ +import sys +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], plugins=[CodeFlashPlugin()] + ) + except Exception as e: + print(f"Failed to collect tests: {e!s}") + exitcode = -1 \ No newline at end of file diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py new file mode 100644 index 000000000..2d3acdd66 --- /dev/null +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -0,0 +1,20 @@ +from __future__ import annotations +from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE +from pathlib import Path +import subprocess + +def trace_benchmarks_pytest(benchmarks_root: Path, project_root: Path, function_list: list[str] = []) -> None: + result = subprocess.run( + [ + SAFE_SYS_EXECUTABLE, + Path(__file__).parent / "pytest_new_process_trace_benchmarks.py", + str(benchmarks_root), + ",".join(function_list) + ], + cwd=project_root, + check=False, + capture_output=True, + text=True, + ) + print("stdout:", result.stdout) + print("stderr:", result.stderr) diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 6ac4db420..04445f1db 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -62,6 +62,10 @@ def parse_args() -> Namespace: ) parser.add_argument("-v", "--verbose", action="store_true", help="Print verbose debug logs") parser.add_argument("--version", action="store_true", help="Print the version of codeflash") + parser.add_argument("--benchmark", action="store_true", help="Trace benchmark tests and calculate optimization impact on benchmarks") + parser.add_argument( + "--benchmarks-root", type=str, help="Path to the directory of the project, where all the pytest-benchmark tests are located." + ) args: Namespace = parser.parse_args() return process_and_validate_cmd_args(args) @@ -116,6 +120,7 @@ def process_pyproject_config(args: Namespace) -> Namespace: "disable_telemetry", "disable_imports_sorting", "git_remote", + "benchmarks_root" ] for key in supported_keys: if key in pyproject_config and ( @@ -127,23 +132,17 @@ def process_pyproject_config(args: Namespace) -> Namespace: assert Path(args.module_root).is_dir(), f"--module-root {args.module_root} must be a valid directory" assert args.tests_root is not None, "--tests-root must be specified" assert Path(args.tests_root).is_dir(), f"--tests-root {args.tests_root} must be a valid directory" - - if env_utils.get_pr_number() is not None: - assert env_utils.ensure_codeflash_api_key(), ( - "Codeflash API key not found. When running in a Github Actions Context, provide the " - "'CODEFLASH_API_KEY' environment variable as a secret.\n" - "You can add a secret by going to your repository's settings page, then clicking 'Secrets' in the left sidebar.\n" - "Then, click 'New repository secret' and add your api key with the variable name CODEFLASH_API_KEY.\n" - f"Here's a direct link: {get_github_secrets_page_url()}\n" - "Exiting..." - ) - - repo = git.Repo(search_parent_directories=True) - - owner, repo_name = get_repo_owner_and_name(repo) - - require_github_app_or_exit(owner, repo_name) - + if args.benchmark: + assert args.benchmarks_root is not None, "--benchmarks-root must be specified when running with --benchmark" + assert Path(args.benchmarks_root).is_dir(), f"--benchmarks-root {args.benchmarks_root} must be a valid directory" + assert not (env_utils.get_pr_number() is not None and not env_utils.ensure_codeflash_api_key()), ( + "Codeflash API key not found. When running in a Github Actions Context, provide the " + "'CODEFLASH_API_KEY' environment variable as a secret.\n" + "You can add a secret by going to your repository's settings page, then clicking 'Secrets' in the left sidebar.\n" + "Then, click 'New repository secret' and add your api key with the variable name CODEFLASH_API_KEY.\n" + f"Here's a direct link: {get_github_secrets_page_url()}\n" + "Exiting..." + ) if hasattr(args, "ignore_paths") and args.ignore_paths is not None: normalized_ignore_paths = [] for path in args.ignore_paths: diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py index 4820d6e65..e26680e1a 100644 --- a/codeflash/discovery/discover_unit_tests.py +++ b/codeflash/discovery/discover_unit_tests.py @@ -107,8 +107,6 @@ def discover_tests_pytest( test_type = TestType.REPLAY_TEST elif "test_concolic_coverage" in test["test_file"]: test_type = TestType.CONCOLIC_COVERAGE_TEST - elif test["test_type"] == "benchmark": # New condition for benchmark tests - test_type = TestType.BENCHMARK_TEST else: test_type = TestType.EXISTING_UNIT_TEST diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index a234a2827..774571de3 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -121,7 +121,6 @@ class FunctionToOptimize: method extends this with the module name from the project root. """ - function_name: str file_path: Path parents: list[FunctionParent] # list[ClassDef | FunctionDef | AsyncFunctionDef] @@ -145,6 +144,11 @@ def qualified_name(self) -> str: def qualified_name_with_modules_from_root(self, project_root_path: Path) -> str: return f"{module_name_from_file_path(self.file_path, project_root_path)}.{self.qualified_name}" + @property + def qualified_name_with_file_name(self) -> str: + class_name = self.parents[0].name if self.parents else None + return f"{self.file_path}:{(class_name + ':' if class_name else '')}{self.function_name}" + def get_functions_to_optimize( optimize_all: str | None, diff --git a/codeflash/discovery/pytest_new_process_discover_benchmarks.py b/codeflash/discovery/pytest_new_process_discover_benchmarks.py new file mode 100644 index 000000000..83175218b --- /dev/null +++ b/codeflash/discovery/pytest_new_process_discover_benchmarks.py @@ -0,0 +1,54 @@ +import sys +from typing import Any + +# This script should not have any relation to the codeflash package, be careful with imports +cwd = sys.argv[1] +tests_root = sys.argv[2] +pickle_path = sys.argv[3] +collected_tests = [] +pytest_rootdir = None +sys.path.insert(1, str(cwd)) + + +class PytestCollectionPlugin: + def pytest_collection_finish(self, session) -> None: + global pytest_rootdir + collected_tests.extend(session.items) + pytest_rootdir = session.config.rootdir + + +def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, str]]: + test_results = [] + for test in pytest_tests: + test_class = None + if test.cls: + test_class = test.parent.name + + # Determine if this is a benchmark test by checking for the benchmark fixture + is_benchmark = hasattr(test, 'fixturenames') and 'benchmark' in test.fixturenames + test_type = 'benchmark' if is_benchmark else 'regular' + + test_results.append({ + "test_file": str(test.path), + "test_class": test_class, + "test_function": test.name, + "test_type": test_type + }) + return test_results + + +if __name__ == "__main__": + import pytest + + try: + exitcode = pytest.main( + [tests_root, "-pno:logging", "--collect-only", "-m", "not skip"], plugins=[PytestCollectionPlugin()] + ) + except Exception as e: + print(f"Failed to collect tests: {e!s}") + exitcode = -1 + tests = parse_pytest_collection_results(collected_tests) + import pickle + + with open(pickle_path, "wb") as f: + pickle.dump((exitcode, tests, pytest_rootdir), f, protocol=pickle.HIGHEST_PROTOCOL) diff --git a/codeflash/discovery/pytest_new_process_discovery.py b/codeflash/discovery/pytest_new_process_discovery.py index fa2dea5e9..2d8583255 100644 --- a/codeflash/discovery/pytest_new_process_discovery.py +++ b/codeflash/discovery/pytest_new_process_discovery.py @@ -29,17 +29,7 @@ def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, s test_class = None if test.cls: test_class = test.parent.name - - # Determine if this is a benchmark test by checking for the benchmark fixture - is_benchmark = hasattr(test, 'fixturenames') and 'benchmark' in test.fixturenames - test_type = 'benchmark' if is_benchmark else 'regular' - - test_results.append({ - "test_file": str(test.path), - "test_class": test_class, - "test_function": test.name, - "test_type": test_type - }) + test_results.append({"test_file": str(test.path), "test_class": test_class, "test_function": test.name}) return test_results diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 93def83c0..66d3c6ab6 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -7,7 +7,7 @@ import subprocess import time import uuid -from collections import defaultdict, deque +from collections import defaultdict from pathlib import Path from typing import TYPE_CHECKING @@ -21,12 +21,12 @@ from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient from codeflash.cli_cmds.console import code_print, console, logger, progress_bar from codeflash.code_utils import env_utils +from codeflash.code_utils.code_extractor import add_needed_imports_from_module, extract_code from codeflash.code_utils.code_replacer import replace_function_definitions_in_module from codeflash.code_utils.code_utils import ( cleanup_paths, file_name_from_test_module_name, get_run_tmp_file, - has_any_async_functions, module_name_from_file_path, ) from codeflash.code_utils.config_consts import ( @@ -37,7 +37,6 @@ ) from codeflash.code_utils.formatter import format_code, sort_imports from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test -from codeflash.code_utils.line_profile_utils import add_decorator_imports from codeflash.code_utils.remove_generated_tests import remove_functions_from_generated_tests from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast from codeflash.code_utils.time_utils import humanize_runtime @@ -49,6 +48,7 @@ BestOptimization, CodeOptimizationContext, FunctionCalledInTest, + FunctionParent, GeneratedTests, GeneratedTestsList, OptimizationSet, @@ -57,9 +57,8 @@ TestFile, TestFiles, TestingMode, - TestResults, - TestType, ) +from codeflash.optimization.function_context import get_constrained_function_context_and_helper_functions from codeflash.result.create_pr import check_create_pr, existing_tests_source_for from codeflash.result.critic import coverage_critic, performance_gain, quantity_of_tests_critic, speedup_critic from codeflash.result.explanation import Explanation @@ -67,15 +66,18 @@ from codeflash.verification.concolic_testing import generate_concolic_tests from codeflash.verification.equivalence import compare_test_results from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture -from codeflash.verification.parse_line_profile_test_output import parse_line_profile_results from codeflash.verification.parse_test_output import parse_test_results -from codeflash.verification.test_runner import run_behavioral_tests, run_benchmarking_tests, run_line_profile_tests +from codeflash.verification.test_results import TestResults, TestType +from codeflash.verification.test_runner import run_behavioral_tests, run_benchmarking_tests from codeflash.verification.verification_utils import get_test_file_path from codeflash.verification.verifier import generate_tests if TYPE_CHECKING: from argparse import Namespace + import numpy as np + import numpy.typing as npt + from codeflash.either import Result from codeflash.models.models import CoverageData, FunctionSource, OptimizedCandidate from codeflash.verification.verification_utils import TestConfig @@ -90,6 +92,8 @@ def __init__( function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_ast: ast.FunctionDef | None = None, aiservice_client: AiServiceClient | None = None, + function_benchmark_timings: dict[str, dict[str, float]] | None = None, + total_benchmark_timings: dict[str, float] | None = None, args: Namespace | None = None, ) -> None: self.project_root = test_cfg.project_root_path @@ -118,6 +122,9 @@ def __init__( self.function_trace_id: str = str(uuid.uuid4()) self.original_module_path = module_name_from_file_path(self.function_to_optimize.file_path, self.project_root) + self.function_benchmark_timings = function_benchmark_timings if function_benchmark_timings else {} + self.total_benchmark_timings = total_benchmark_timings if total_benchmark_timings else {} + def optimize_function(self) -> Result[BestOptimization, str]: should_run_experiment = self.experiment_id is not None logger.debug(f"Function Trace ID: {self.function_trace_id}") @@ -134,10 +141,19 @@ def optimize_function(self) -> Result[BestOptimization, str]: with helper_function_path.open(encoding="utf8") as f: helper_code = f.read() original_helper_code[helper_function_path] = helper_code - if has_any_async_functions(code_context.read_writable_code): - return Failure("Codeflash does not support async functions in the code to optimize.") + + logger.info("Code to be optimized:") code_print(code_context.read_writable_code) + for module_abspath, helper_code_source in original_helper_code.items(): + code_context.code_to_optimize_with_helpers = add_needed_imports_from_module( + helper_code_source, + code_context.code_to_optimize_with_helpers, + module_abspath, + self.function_to_optimize.file_path, + self.args.project_root, + ) + generated_test_paths = [ get_test_file_path( self.test_cfg.tests_root, self.function_to_optimize.function_name, test_index, test_type="unit" @@ -156,7 +172,7 @@ def optimize_function(self) -> Result[BestOptimization, str]: transient=True, ): generated_results = self.generate_tests_and_optimizations( - testgen_context_code=code_context.testgen_context_code, + code_to_optimize_with_helpers=code_context.code_to_optimize_with_helpers, read_writable_code=code_context.read_writable_code, read_only_context_code=code_context.read_only_context_code, helper_functions=code_context.helper_functions, @@ -232,11 +248,10 @@ def optimize_function(self) -> Result[BestOptimization, str]: ): cleanup_paths(paths_to_cleanup) return Failure("The threshold for test coverage was not met.") - # request for new optimizations but don't block execution, check for completion later - # adding to control and experiment set but with same traceid + best_optimization = None - for _u, candidates in enumerate([optimizations_set.control, optimizations_set.experiment]): + for u, candidates in enumerate([optimizations_set.control, optimizations_set.experiment]): if candidates is None: continue @@ -270,6 +285,20 @@ def optimize_function(self) -> Result[BestOptimization, str]: function_name=function_to_optimize_qualified_name, file_path=self.function_to_optimize.file_path, ) + speedup = explanation.speedup # eg. 1.2 means 1.2x faster + if self.args.benchmark: + fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_file_name] + for benchmark_name, og_benchmark_timing in fto_benchmark_timings.items(): + print(f"Calculating speedup for benchmark {benchmark_name}") + total_benchmark_timing = self.total_benchmark_timings[benchmark_name] + # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values + expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + og_benchmark_timing / speedup + print(f"Expected new benchmark timing: {expected_new_benchmark_timing}") + print(f"Original benchmark timing: {total_benchmark_timing}") + print(f"Benchmark speedup: {total_benchmark_timing / expected_new_benchmark_timing}") + + speedup = total_benchmark_timing / expected_new_benchmark_timing + print(f"Speedup: {speedup}") self.log_successful_optimization(explanation, generated_tests) @@ -359,123 +388,94 @@ def determine_best_candidate( f"{self.function_to_optimize.qualified_name}…" ) console.rule() - candidates = deque(candidates) - # Start a new thread for AI service request, start loop in main thread - # check if aiservice request is complete, when it is complete, append result to the candidates list - with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: - future_line_profile_results = executor.submit( - self.aiservice_client.optimize_python_code_line_profiler, - source_code=code_context.read_writable_code, - dependency_code=code_context.read_only_context_code, - trace_id=self.function_trace_id, - line_profiler_results=original_code_baseline.line_profile_results["str_out"], - num_candidates=10, - experiment_metadata=None, - ) - try: - candidate_index = 0 - done = False - original_len = len(candidates) - while candidates: - # for candidate_index, candidate in enumerate(candidates, start=1): - done = True if future_line_profile_results is None else future_line_profile_results.done() - if done and (future_line_profile_results is not None): - line_profile_results = future_line_profile_results.result() - candidates.extend(line_profile_results) - original_len+= len(line_profile_results) - logger.info(f"Added {len(line_profile_results)} results from line profiler to candidates, total candidates now: {original_len}") - future_line_profile_results = None - candidate_index += 1 - candidate = candidates.popleft() - get_run_tmp_file(Path(f"test_return_values_{candidate_index}.bin")).unlink(missing_ok=True) - get_run_tmp_file(Path(f"test_return_values_{candidate_index}.sqlite")).unlink(missing_ok=True) - logger.info(f"Optimization candidate {candidate_index}/{original_len}:") - code_print(candidate.source_code) - try: - did_update = self.replace_function_and_helpers_with_optimized_code( - code_context=code_context, optimized_code=candidate.source_code - ) - if not did_update: - logger.warning( - "No functions were replaced in the optimized code. Skipping optimization candidate." - ) - console.rule() - continue - except (ValueError, SyntaxError, cst.ParserSyntaxError, AttributeError) as e: - logger.error(e) - self.write_code_and_helpers( - self.function_to_optimize_source_code, - original_helper_code, - self.function_to_optimize.file_path, - ) - continue - - # Instrument codeflash capture - run_results = self.run_optimized_candidate( - optimization_candidate_index=candidate_index, - baseline_results=original_code_baseline, - original_helper_code=original_helper_code, - file_path_to_helper_classes=file_path_to_helper_classes, + try: + for candidate_index, candidate in enumerate(candidates, start=1): + get_run_tmp_file(Path(f"test_return_values_{candidate_index}.bin")).unlink(missing_ok=True) + get_run_tmp_file(Path(f"test_return_values_{candidate_index}.sqlite")).unlink(missing_ok=True) + logger.info(f"Optimization candidate {candidate_index}/{len(candidates)}:") + code_print(candidate.source_code) + try: + did_update = self.replace_function_and_helpers_with_optimized_code( + code_context=code_context, optimized_code=candidate.source_code ) - console.rule() - - if not is_successful(run_results): - optimized_runtimes[candidate.optimization_id] = None - is_correct[candidate.optimization_id] = False - speedup_ratios[candidate.optimization_id] = None - else: - candidate_result: OptimizedCandidateResult = run_results.unwrap() - best_test_runtime = candidate_result.best_test_runtime - optimized_runtimes[candidate.optimization_id] = best_test_runtime - is_correct[candidate.optimization_id] = True - perf_gain = performance_gain( - original_runtime_ns=original_code_baseline.runtime, optimized_runtime_ns=best_test_runtime + if not did_update: + logger.warning( + "No functions were replaced in the optimized code. Skipping optimization candidate." ) - speedup_ratios[candidate.optimization_id] = perf_gain - - tree = Tree(f"Candidate #{candidate_index} - Runtime Information") - if speedup_critic( - candidate_result, original_code_baseline.runtime, best_runtime_until_now - ) and quantity_of_tests_critic(candidate_result): - tree.add("This candidate is faster than the previous best candidate. πŸš€") - tree.add(f"Original summed runtime: {humanize_runtime(original_code_baseline.runtime)}") - tree.add( - f"Best summed runtime: {humanize_runtime(candidate_result.best_test_runtime)} " - f"(measured over {candidate_result.max_loop_count} " - f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" - ) - tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") - - best_optimization = BestOptimization( - candidate=candidate, - helper_functions=code_context.helper_functions, - runtime=best_test_runtime, - winning_behavioral_test_results=candidate_result.behavior_test_results, - winning_benchmarking_test_results=candidate_result.benchmarking_test_results, - ) - best_runtime_until_now = best_test_runtime - else: - tree.add( - f"Summed runtime: {humanize_runtime(best_test_runtime)} " - f"(measured over {candidate_result.max_loop_count} " - f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" - ) - tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X") - console.print(tree) console.rule() - + continue + except (ValueError, SyntaxError, cst.ParserSyntaxError, AttributeError) as e: + logger.error(e) self.write_code_and_helpers( self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path ) + continue + + # Instrument codeflash capture + run_results = self.run_optimized_candidate( + optimization_candidate_index=candidate_index, + baseline_results=original_code_baseline, + original_helper_code=original_helper_code, + file_path_to_helper_classes=file_path_to_helper_classes, + ) + console.rule() + + if not is_successful(run_results): + optimized_runtimes[candidate.optimization_id] = None + is_correct[candidate.optimization_id] = False + speedup_ratios[candidate.optimization_id] = None + else: + candidate_result: OptimizedCandidateResult = run_results.unwrap() + best_test_runtime = candidate_result.best_test_runtime + optimized_runtimes[candidate.optimization_id] = best_test_runtime + is_correct[candidate.optimization_id] = True + perf_gain = performance_gain( + original_runtime_ns=original_code_baseline.runtime, optimized_runtime_ns=best_test_runtime + ) + speedup_ratios[candidate.optimization_id] = perf_gain + + tree = Tree(f"Candidate #{candidate_index} - Runtime Information") + if speedup_critic( + candidate_result, original_code_baseline.runtime, best_runtime_until_now + ) and quantity_of_tests_critic(candidate_result): + tree.add("This candidate is faster than the previous best candidate. πŸš€") + tree.add(f"Original summed runtime: {humanize_runtime(original_code_baseline.runtime)}") + tree.add( + f"Best summed runtime: {humanize_runtime(candidate_result.best_test_runtime)} " + f"(measured over {candidate_result.max_loop_count} " + f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + ) + tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") + + best_optimization = BestOptimization( + candidate=candidate, + helper_functions=code_context.helper_functions, + runtime=best_test_runtime, + winning_behavioral_test_results=candidate_result.behavior_test_results, + winning_benchmarking_test_results=candidate_result.benchmarking_test_results, + ) + best_runtime_until_now = best_test_runtime + else: + tree.add( + f"Summed runtime: {humanize_runtime(best_test_runtime)} " + f"(measured over {candidate_result.max_loop_count} " + f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + ) + tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X") + console.print(tree) + console.rule() - except KeyboardInterrupt as e: self.write_code_and_helpers( self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path ) - logger.exception(f"Optimization interrupted: {e}") - raise + except KeyboardInterrupt as e: + self.write_code_and_helpers( + self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path + ) + logger.exception(f"Optimization interrupted: {e}") + raise self.aiservice_client.log_results( function_trace_id=self.function_trace_id, @@ -575,6 +575,50 @@ def replace_function_and_helpers_with_optimized_code( return did_update def get_code_optimization_context(self) -> Result[CodeOptimizationContext, str]: + code_to_optimize, contextual_dunder_methods = extract_code([self.function_to_optimize]) + if code_to_optimize is None: + return Failure("Could not find function to optimize.") + (helper_code, helper_functions, helper_dunder_methods) = get_constrained_function_context_and_helper_functions( + self.function_to_optimize, self.project_root, code_to_optimize + ) + if self.function_to_optimize.parents: + function_class = self.function_to_optimize.parents[0].name + same_class_helper_methods = [ + df + for df in helper_functions + if df.qualified_name.count(".") > 0 and df.qualified_name.split(".")[0] == function_class + ] + optimizable_methods = [ + FunctionToOptimize( + df.qualified_name.split(".")[-1], + df.file_path, + [FunctionParent(df.qualified_name.split(".")[0], "ClassDef")], + None, + None, + ) + for df in same_class_helper_methods + ] + [self.function_to_optimize] + dedup_optimizable_methods = [] + added_methods = set() + for method in reversed(optimizable_methods): + if f"{method.file_path}.{method.qualified_name}" not in added_methods: + dedup_optimizable_methods.append(method) + added_methods.add(f"{method.file_path}.{method.qualified_name}") + if len(dedup_optimizable_methods) > 1: + code_to_optimize, contextual_dunder_methods = extract_code(list(reversed(dedup_optimizable_methods))) + if code_to_optimize is None: + return Failure("Could not find function to optimize.") + code_to_optimize_with_helpers = helper_code + "\n" + code_to_optimize + + code_to_optimize_with_helpers_and_imports = add_needed_imports_from_module( + self.function_to_optimize_source_code, + code_to_optimize_with_helpers, + self.function_to_optimize.file_path, + self.function_to_optimize.file_path, + self.project_root, + helper_functions, + ) + try: new_code_ctx = code_context_extractor.get_code_optimization_context( self.function_to_optimize, self.project_root @@ -584,7 +628,7 @@ def get_code_optimization_context(self) -> Result[CodeOptimizationContext, str]: return Success( CodeOptimizationContext( - testgen_context_code=new_code_ctx.testgen_context_code, + code_to_optimize_with_helpers=code_to_optimize_with_helpers_and_imports, read_writable_code=new_code_ctx.read_writable_code, read_only_context_code=new_code_ctx.read_only_context_code, helper_functions=new_code_ctx.helper_functions, # only functions that are read writable @@ -686,7 +730,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi def generate_tests_and_optimizations( self, - testgen_context_code: str, + code_to_optimize_with_helpers: str, read_writable_code: str, read_only_context_code: str, helper_functions: list[FunctionSource], @@ -701,7 +745,7 @@ def generate_tests_and_optimizations( # Submit the test generation task as future future_tests = self.generate_and_instrument_tests( executor, - testgen_context_code, + code_to_optimize_with_helpers, [definition.fully_qualified_name for definition in helper_functions], generated_test_paths, generated_perf_test_paths, @@ -790,7 +834,6 @@ def establish_original_code_baseline( original_helper_code: dict[Path, str], file_path_to_helper_classes: dict[Path, set[str]], ) -> Result[tuple[OriginalCodeBaseline, list[str]], str]: - line_profile_results = {"timings": {}, "unit": 0, "str_out": ""} # For the original function - run the tests and get the runtime, plus coverage with progress_bar(f"Establishing original code baseline for {self.function_to_optimize.function_name}"): assert (test_framework := self.args.test_framework) in ["pytest", "unittest"] @@ -831,31 +874,11 @@ def establish_original_code_baseline( ) console.rule() return Failure("Failed to establish a baseline for the original code - bevhavioral tests failed.") - if not coverage_critic(coverage_results, self.args.test_framework): + if not coverage_critic( + coverage_results, self.args.test_framework + ): return Failure("The threshold for test coverage was not met.") if test_framework == "pytest": - try: - line_profiler_output_file = add_decorator_imports(self.function_to_optimize, code_context) - line_profile_results, _ = self.run_and_parse_tests( - testing_type=TestingMode.LINE_PROFILE, - test_env=test_env, - test_files=self.test_files, - optimization_iteration=0, - testing_time=TOTAL_LOOPING_TIME, - enable_coverage=False, - code_context=code_context, - line_profiler_output_file=line_profiler_output_file, - ) - finally: - # Remove codeflash capture - self.write_code_and_helpers( - self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path - ) - if line_profile_results["str_out"] == "": - logger.warning( - f"Couldn't run line profiler for original function {self.function_to_optimize.function_name}" - ) - console.rule() benchmarking_results, _ = self.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, @@ -894,6 +917,7 @@ def establish_original_code_baseline( ) console.rule() + total_timing = benchmarking_results.total_passed_runtime() # caution: doesn't handle the loop index functions_to_remove = [ result.id.test_function_name @@ -927,7 +951,6 @@ def establish_original_code_baseline( benchmarking_test_results=benchmarking_results, runtime=total_timing, coverage_results=coverage_results, - line_profile_results=line_profile_results, ), functions_to_remove, ) @@ -1063,77 +1086,59 @@ def run_and_parse_tests( pytest_max_loops: int = 100_000, code_context: CodeOptimizationContext | None = None, unittest_loop_index: int | None = None, - line_profiler_output_file: Path | None = None, - ) -> tuple[TestResults | dict, CoverageData | None]: + ) -> tuple[TestResults, CoverageData | None]: coverage_database_file = None - coverage_config_file = None try: if testing_type == TestingMode.BEHAVIOR: - result_file_path, run_result, coverage_database_file, coverage_config_file = run_behavioral_tests( + result_file_path, run_result, coverage_database_file = run_behavioral_tests( test_files, test_framework=self.test_cfg.test_framework, cwd=self.project_root, test_env=test_env, pytest_timeout=INDIVIDUAL_TESTCASE_TIMEOUT, + pytest_cmd=self.test_cfg.pytest_cmd, verbose=True, enable_coverage=enable_coverage, ) - elif testing_type == TestingMode.LINE_PROFILE: - result_file_path, run_result = run_line_profile_tests( - test_files, - cwd=self.project_root, - test_env=test_env, - pytest_cmd=self.test_cfg.pytest_cmd, - pytest_timeout=INDIVIDUAL_TESTCASE_TIMEOUT, - pytest_target_runtime_seconds=testing_time, - pytest_min_loops=pytest_min_loops, - pytest_max_loops=pytest_min_loops, - test_framework=self.test_cfg.test_framework, - line_profiler_output_file=line_profiler_output_file, - ) elif testing_type == TestingMode.PERFORMANCE: result_file_path, run_result = run_benchmarking_tests( test_files, cwd=self.project_root, test_env=test_env, - pytest_cmd=self.test_cfg.pytest_cmd, pytest_timeout=INDIVIDUAL_TESTCASE_TIMEOUT, + pytest_cmd=self.test_cfg.pytest_cmd, pytest_target_runtime_seconds=testing_time, pytest_min_loops=pytest_min_loops, pytest_max_loops=pytest_max_loops, test_framework=self.test_cfg.test_framework, ) else: - msg = f"Unexpected testing type: {testing_type}" - raise ValueError(msg) + raise ValueError(f"Unexpected testing type: {testing_type}") except subprocess.TimeoutExpired: logger.exception( - f"Error running tests in {', '.join(str(f) for f in test_files.test_files)}.\nTimeout Error" + f'Error running tests in {", ".join(str(f) for f in test_files.test_files)}.\nTimeout Error' ) return TestResults(), None if run_result.returncode != 0 and testing_type == TestingMode.BEHAVIOR: logger.debug( - f"Nonzero return code {run_result.returncode} when running tests in " - f"{', '.join([str(f.instrumented_behavior_file_path) for f in test_files.test_files])}.\n" + f'Nonzero return code {run_result.returncode} when running tests in ' + f'{", ".join([str(f.instrumented_behavior_file_path) for f in test_files.test_files])}.\n' f"stdout: {run_result.stdout}\n" f"stderr: {run_result.stderr}\n" ) - if testing_type in [TestingMode.BEHAVIOR, TestingMode.PERFORMANCE]: - results, coverage_results = parse_test_results( - test_xml_path=result_file_path, - test_files=test_files, - test_config=self.test_cfg, - optimization_iteration=optimization_iteration, - run_result=run_result, - unittest_loop_index=unittest_loop_index, - function_name=self.function_to_optimize.function_name, - source_file=self.function_to_optimize.file_path, - code_context=code_context, - coverage_database_file=coverage_database_file, - coverage_config_file=coverage_config_file, - ) - else: - results, coverage_results = parse_line_profile_results(line_profiler_output_file=line_profiler_output_file) + # print(test_files) + results, coverage_results = parse_test_results( + test_xml_path=result_file_path, + test_files=test_files, + test_config=self.test_cfg, + optimization_iteration=optimization_iteration, + run_result=run_result, + unittest_loop_index=unittest_loop_index, + function_name=self.function_to_optimize.function_name, + source_file=self.function_to_optimize.file_path, + code_context=code_context, + coverage_database_file=coverage_database_file, + ) return results, coverage_results def generate_and_instrument_tests( @@ -1163,3 +1168,4 @@ def generate_and_instrument_tests( zip(generated_test_paths, generated_perf_test_paths) ) ] + diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 8eae1014a..01a196143 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -8,7 +8,8 @@ from typing import TYPE_CHECKING from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient -from codeflash.cli_cmds.console import console, logger, progress_bar +from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest +from codeflash.cli_cmds.console import console, logger from codeflash.code_utils import env_utils from codeflash.code_utils.code_replacer import normalize_code, normalize_node from codeflash.code_utils.code_utils import get_run_tmp_file @@ -16,10 +17,12 @@ from codeflash.discovery.discover_unit_tests import discover_unit_tests from codeflash.discovery.functions_to_optimize import get_functions_to_optimize from codeflash.either import is_successful -from codeflash.models.models import TestType, ValidCode +from codeflash.models.models import TestFiles, ValidCode from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.telemetry.posthog_cf import ph +from codeflash.verification.test_results import TestType from codeflash.verification.verification_utils import TestConfig +from codeflash.benchmarking.get_trace_info import get_function_benchmark_timings, get_benchmark_timings if TYPE_CHECKING: from argparse import Namespace @@ -50,6 +53,8 @@ def create_function_optimizer( function_to_optimize_ast: ast.FunctionDef | None = None, function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_source_code: str | None = "", + function_benchmark_timings: dict[str, dict[str, float]] | None = None, + total_benchmark_timings: dict[str, float] | None = None, ) -> FunctionOptimizer: return FunctionOptimizer( function_to_optimize=function_to_optimize, @@ -59,6 +64,8 @@ def create_function_optimizer( function_to_optimize_ast=function_to_optimize_ast, aiservice_client=self.aiservice_client, args=self.args, + function_benchmark_timings=function_benchmark_timings if function_benchmark_timings else None, + total_benchmark_timings=total_benchmark_timings if total_benchmark_timings else None, ) def run(self) -> None: @@ -80,6 +87,23 @@ def run(self) -> None: project_root=self.args.project_root, 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) + print(function_benchmark_timings) + total_benchmark_timings = get_benchmark_timings(trace_dir) + print("Total benchmark timings:") + print(total_benchmark_timings) + # for function in fully_qualified_function_names: + optimizations_found: int = 0 function_iterator_count: int = 0 @@ -93,6 +117,8 @@ def run(self) -> None: logger.info("No functions found to optimize. Exiting…") return + console.rule() + logger.info(f"Discovering existing unit tests in {self.test_cfg.tests_root}…") console.rule() function_to_tests: dict[str, list[FunctionCalledInTest]] = discover_unit_tests(self.test_cfg) num_discovered_tests: int = sum([len(value) for value in function_to_tests.values()]) @@ -136,7 +162,6 @@ def run(self) -> None: validated_original_code[analysis.file_path] = ValidCode( source_code=callee_original_code, normalized_code=normalized_callee_original_code ) - if has_syntax_error: continue @@ -146,7 +171,7 @@ def run(self) -> None: f"Optimizing function {function_iterator_count} of {num_optimizable_functions}: " f"{function_to_optimize.qualified_name}" ) - console.rule() + if not ( function_to_optimize_ast := get_first_top_level_function_or_method_ast( function_to_optimize.function_name, function_to_optimize.parents, original_module_ast @@ -157,12 +182,17 @@ def run(self) -> None: f"Skipping optimization." ) continue - function_optimizer = self.create_function_optimizer( - function_to_optimize, - function_to_optimize_ast, - function_to_tests, - validated_original_code[original_module_path].source_code, - ) + if self.args.benchmark: + + function_optimizer = self.create_function_optimizer( + function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code, function_benchmark_timings, total_benchmark_timings + ) + else: + function_optimizer = self.create_function_optimizer( + function_to_optimize, function_to_optimize_ast, function_to_tests, + validated_original_code[original_module_path].source_code + ) + best_optimization = function_optimizer.optimize_function() if is_successful(best_optimization): optimizations_found += 1 @@ -191,6 +221,7 @@ def run(self) -> None: get_run_tmp_file.tmpdir.cleanup() + def run_with_args(args: Namespace) -> None: optimizer = Optimizer(args) optimizer.run() diff --git a/codeflash/tracer.py b/codeflash/tracer.py index 5d1240868..5bc1ae482 100644 --- a/codeflash/tracer.py +++ b/codeflash/tracer.py @@ -18,21 +18,19 @@ import os import pathlib import pickle +import re import sqlite3 import sys -import threading import time -from argparse import ArgumentParser from collections import defaultdict +from copy import copy +from io import StringIO from pathlib import Path -from typing import TYPE_CHECKING, Any, Callable, ClassVar +from types import FrameType +from typing import Any, ClassVar, List import dill import isort -from rich.align import Align -from rich.panel import Panel -from rich.table import Table -from rich.text import Text from codeflash.cli_cmds.cli import project_root_from_module_root from codeflash.cli_cmds.console import console @@ -42,34 +40,14 @@ from codeflash.tracing.replay_test import create_trace_replay_test from codeflash.tracing.tracing_utils import FunctionModules from codeflash.verification.verification_utils import get_test_file_path - -if TYPE_CHECKING: - from types import FrameType, TracebackType - - -class FakeCode: - def __init__(self, filename: str, line: int, name: str) -> None: - self.co_filename = filename - self.co_line = line - self.co_name = name - self.co_firstlineno = 0 - - def __repr__(self) -> str: - return repr((self.co_filename, self.co_line, self.co_name, None)) - - -class FakeFrame: - def __init__(self, code: FakeCode, prior: FakeFrame | None) -> None: - self.f_code = code - self.f_back = prior - self.f_locals: dict = {} - +# import warnings +# warnings.filterwarnings("ignore", category=dill.PickleWarning) +# warnings.filterwarnings("ignore", category=DeprecationWarning) # Debug this file by simply adding print statements. This file is not meant to be debugged by the debugger. class Tracer: - """Use this class as a 'with' context manager to trace a function call. - - Traces function calls, input arguments, and profiling info. + """Use this class as a 'with' context manager to trace a function call, + input arguments, and profiling info. """ def __init__( @@ -81,9 +59,7 @@ def __init__( max_function_count: int = 256, timeout: int | None = None, # seconds ) -> None: - """Use this class to trace function calls. - - :param output: The path to the output trace file + """:param output: The path to the output trace file :param functions: List of functions to trace. If None, trace all functions :param disable: Disable the tracer if True :param config_file_path: Path to the pyproject.toml file, if None then it will be auto-discovered @@ -94,9 +70,7 @@ def __init__( if functions is None: functions = [] if os.environ.get("CODEFLASH_TRACER_DISABLE", "0") == "1": - console.rule( - "Codeflash: Tracer disabled by environment variable CODEFLASH_TRACER_DISABLE", style="bold red" - ) + console.print("Codeflash: Tracer disabled by environment variable CODEFLASH_TRACER_DISABLE") disable = True self.disable = disable if self.disable: @@ -111,7 +85,7 @@ def __init__( self.con = None self.output_file = Path(output).resolve() self.functions = functions - self.function_modules: list[FunctionModules] = [] + self.function_modules: List[FunctionModules] = [] self.function_count = defaultdict(int) self.current_file_path = Path(__file__).resolve() self.ignored_qualified_functions = { @@ -121,10 +95,10 @@ def __init__( self.max_function_count = max_function_count self.config, found_config_path = parse_config_file(config_file_path) self.project_root = project_root_from_module_root(Path(self.config["module_root"]), found_config_path) - console.rule(f"Project Root: {self.project_root}", style="bold blue") + print("project_root", self.project_root) self.ignored_functions = {"", "", "", "", "", ""} - self.file_being_called_from: str = str(Path(sys._getframe().f_back.f_code.co_filename).name).replace(".", "_") # noqa: SLF001 + self.file_being_called_from: str = str(Path(sys._getframe().f_back.f_code.co_filename).name).replace(".", "_") assert timeout is None or timeout > 0, "Timeout should be greater than 0" self.timeout = timeout @@ -145,44 +119,48 @@ def __init__( def __enter__(self) -> None: if self.disable: return - if getattr(Tracer, "used_once", False): - console.print( - "Codeflash: Tracer can only be used once per program run. " - "Please only enable the Tracer once. Skipping tracing this section." - ) - self.disable = True - return - Tracer.used_once = True + + # if getattr(Tracer, "used_once", False): + # console.print( + # "Codeflash: Tracer can only be used once per program run. " + # "Please only enable the Tracer once. Skipping tracing this section." + # ) + # self.disable = True + # return + # Tracer.used_once = True if pathlib.Path(self.output_file).exists(): - console.rule("Removing existing trace file", style="bold red") - console.rule() + console.print("Codeflash: Removing existing trace file") pathlib.Path(self.output_file).unlink(missing_ok=True) - self.con = sqlite3.connect(self.output_file, check_same_thread=False) + self.con = sqlite3.connect(self.output_file) cur = self.con.cursor() cur.execute("""PRAGMA synchronous = OFF""") - cur.execute("""PRAGMA journal_mode = WAL""") # TODO: Check out if we need to export the function test name as well cur.execute( "CREATE TABLE function_calls(type TEXT, function TEXT, classname TEXT, filename TEXT, " "line_number INTEGER, last_frame_address INTEGER, time_ns INTEGER, args BLOB)" ) - console.rule("Codeflash: Traced Program Output Begin", style="bold blue") - frame = sys._getframe(0) # Get this frame and simulate a call to it # noqa: SLF001 + console.print("Codeflash: Tracing started!") + frame = sys._getframe(0) # Get this frame and simulate a call to it self.dispatch["call"](self, frame, 0) self.start_time = time.time() sys.setprofile(self.trace_callback) - threading.setprofile(self.trace_callback) - def __exit__( - self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None - ) -> None: + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: if self.disable: return sys.setprofile(None) self.con.commit() - console.rule("Codeflash: Traced Program Output End", style="bold blue") + # Check if any functions were actually traced + if self.trace_count == 0: + self.con.close() + # Delete the trace file if no functions were traced + if self.output_file.exists(): + self.output_file.unlink() + console.print("Codeflash: No functions were traced. Removing trace database.") + return + self.create_stats() cur = self.con.cursor() @@ -226,13 +204,14 @@ def __exit__( test_framework=self.config["test_framework"], max_run_count=self.max_function_count, ) - function_path = "_".join(self.functions) if self.functions else self.file_being_called_from + # Need a better way to store the replay test + # function_path = "_".join(self.functions) if self.functions else self.file_being_called_from + function_path = self.file_being_called_from test_file_path = get_test_file_path( test_dir=Path(self.config["tests_root"]), function_name=function_path, test_type="replay" ) replay_test = isort.code(replay_test) - - with Path(test_file_path).open("w", encoding="utf8") as file: + with open(test_file_path, "w", encoding="utf8") as file: file.write(replay_test) console.print( @@ -242,27 +221,25 @@ def __exit__( overflow="ignore", ) - def tracer_logic(self, frame: FrameType, event: str) -> None: + def tracer_logic(self, frame: FrameType, event: str): if event != "call": return - if self.timeout is not None and (time.time() - self.start_time) > self.timeout: - sys.setprofile(None) - threading.setprofile(None) - console.print(f"Codeflash: Timeout reached! Stopping tracing at {self.timeout} seconds.") - return + if self.timeout is not None: + if (time.time() - self.start_time) > self.timeout: + sys.setprofile(None) + console.print(f"Codeflash: Timeout reached! Stopping tracing at {self.timeout} seconds.") + return code = frame.f_code - file_name = Path(code.co_filename).resolve() # TODO : It currently doesn't log the last return call from the first function if code.co_name in self.ignored_functions: return - if not file_name.is_relative_to(self.project_root): - return if not file_name.exists(): return - if self.functions and code.co_name not in self.functions: - return + # if self.functions: + # if code.co_name not in self.functions: + # return class_name = None arguments = frame.f_locals try: @@ -274,12 +251,16 @@ def tracer_logic(self, frame: FrameType, event: str) -> None: class_name = arguments["self"].__class__.__name__ elif "cls" in arguments and hasattr(arguments["cls"], "__name__"): class_name = arguments["cls"].__name__ - except: # noqa: E722 + except: # someone can override the getattr method and raise an exception. I'm looking at you wrapt return + function_qualified_name = f"{file_name}:{(class_name + ':' if class_name else '')}{code.co_name}" if function_qualified_name in self.ignored_qualified_functions: return + if self.functions and function_qualified_name not in self.functions: + return + if function_qualified_name not in self.function_count: # seeing this function for the first time self.function_count[function_qualified_name] = 0 @@ -354,14 +335,17 @@ def tracer_logic(self, frame: FrameType, event: str) -> None: self.next_insert = 1000 self.con.commit() - def trace_callback(self, frame: FrameType, event: str, arg: str | None) -> None: + def trace_callback(self, frame: FrameType, event: str, arg: Any) -> None: # profiler section timer = self.timer t = timer() - self.t - self.bias if event == "c_call": self.c_func_name = arg.__name__ - prof_success = bool(self.dispatch[event](self, frame, t)) + if self.dispatch[event](self, frame, t): + prof_success = True + else: + prof_success = False # tracer section self.tracer_logic(frame, event) # measure the time as the last thing before return @@ -370,60 +354,45 @@ def trace_callback(self, frame: FrameType, event: str, arg: str | None) -> None: else: self.t = timer() - t # put back unrecorded delta - def trace_dispatch_call(self, frame: FrameType, t: int) -> int: - """Handle call events in the profiler.""" + def trace_dispatch_call(self, frame, t): + if self.cur and frame.f_back is not self.cur[-2]: + rpt, rit, ret, rfn, rframe, rcur = self.cur + if not isinstance(rframe, Tracer.fake_frame): + assert rframe.f_back is frame.f_back, ("Bad call", rfn, rframe, rframe.f_back, frame, frame.f_back) + self.trace_dispatch_return(rframe, 0) + assert self.cur is None or frame.f_back is self.cur[-2], ("Bad call", self.cur[-3]) + fcode = frame.f_code + arguments = frame.f_locals + class_name = None try: - # In multi-threaded contexts, we need to be more careful about frame comparisons - if self.cur and frame.f_back is not self.cur[-2]: - # This happens when we're in a different thread - rpt, rit, ret, rfn, rframe, rcur = self.cur - - # Only attempt to handle the frame mismatch if we have a valid rframe - if ( - not isinstance(rframe, FakeFrame) - and hasattr(rframe, "f_back") - and hasattr(frame, "f_back") - and rframe.f_back is frame.f_back - ): - self.trace_dispatch_return(rframe, 0) - - # Get function information - fcode = frame.f_code - arguments = frame.f_locals - class_name = None - try: - if ( - "self" in arguments - and hasattr(arguments["self"], "__class__") - and hasattr(arguments["self"].__class__, "__name__") - ): - class_name = arguments["self"].__class__.__name__ - elif "cls" in arguments and hasattr(arguments["cls"], "__name__"): - class_name = arguments["cls"].__name__ - except Exception: # noqa: BLE001, S110 - pass - - fn = (fcode.co_filename, fcode.co_firstlineno, fcode.co_name, class_name) - self.cur = (t, 0, 0, fn, frame, self.cur) - timings = self.timings - if fn in timings: - cc, ns, tt, ct, callers = timings[fn] - timings[fn] = cc, ns + 1, tt, ct, callers - else: - timings[fn] = 0, 0, 0, 0, {} - return 1 # noqa: TRY300 - except Exception: # noqa: BLE001 - # Handle any errors gracefully - return 0 + if ( + "self" in arguments + and hasattr(arguments["self"], "__class__") + and hasattr(arguments["self"].__class__, "__name__") + ): + class_name = arguments["self"].__class__.__name__ + elif "cls" in arguments and hasattr(arguments["cls"], "__name__"): + class_name = arguments["cls"].__name__ + except: + pass + fn = (fcode.co_filename, fcode.co_firstlineno, fcode.co_name, class_name) + self.cur = (t, 0, 0, fn, frame, self.cur) + timings = self.timings + if fn in timings: + cc, ns, tt, ct, callers = timings[fn] + timings[fn] = cc, ns + 1, tt, ct, callers + else: + timings[fn] = 0, 0, 0, 0, {} + return 1 - def trace_dispatch_exception(self, frame: FrameType, t: int) -> int: + def trace_dispatch_exception(self, frame, t): rpt, rit, ret, rfn, rframe, rcur = self.cur if (rframe is not frame) and rcur: return self.trace_dispatch_return(rframe, t) self.cur = rpt, rit + t, ret, rfn, rframe, rcur return 1 - def trace_dispatch_c_call(self, frame: FrameType, t: int) -> int: + def trace_dispatch_c_call(self, frame, t): fn = ("", 0, self.c_func_name, None) self.cur = (t, 0, 0, fn, frame, self.cur) timings = self.timings @@ -434,27 +403,15 @@ def trace_dispatch_c_call(self, frame: FrameType, t: int) -> int: timings[fn] = 0, 0, 0, 0, {} return 1 - def trace_dispatch_return(self, frame: FrameType, t: int) -> int: - if not self.cur or not self.cur[-2]: - return 0 - - # In multi-threaded environments, frames can get mismatched + def trace_dispatch_return(self, frame, t): if frame is not self.cur[-2]: - # Don't assert in threaded environments - frames can legitimately differ - if hasattr(frame, "f_back") and hasattr(self.cur[-2], "f_back") and frame is self.cur[-2].f_back: - self.trace_dispatch_return(self.cur[-2], 0) - else: - # We're in a different thread or context, can't continue with this frame - return 0 + assert frame is self.cur[-2].f_back, ("Bad return", self.cur[-3]) + self.trace_dispatch_return(self.cur[-2], 0) + # Prefix "r" means part of the Returning or exiting frame. # Prefix "p" means part of the Previous or Parent or older frame. rpt, rit, ret, rfn, frame, rcur = self.cur - - # Guard against invalid rcur (w threading) - if not rcur: - return 0 - rit = rit + t frame_total = rit + ret @@ -462,9 +419,6 @@ def trace_dispatch_return(self, frame: FrameType, t: int) -> int: self.cur = ppt, pit + rpt, pet + frame_total, pfn, pframe, pcur timings = self.timings - if rfn not in timings: - # w threading, rfn can be missing - timings[rfn] = 0, 0, 0, 0, {} cc, ns, tt, ct, callers = timings[rfn] if not ns: # This is the only occurrence of the function on the stack. @@ -486,7 +440,7 @@ def trace_dispatch_return(self, frame: FrameType, t: int) -> int: return 1 - dispatch: ClassVar[dict[str, Callable[[Tracer, FrameType, int], int]]] = { + dispatch: ClassVar[dict[str, callable]] = { "call": trace_dispatch_call, "exception": trace_dispatch_exception, "return": trace_dispatch_return, @@ -495,13 +449,32 @@ def trace_dispatch_return(self, frame: FrameType, t: int) -> int: "c_return": trace_dispatch_return, } - def simulate_call(self, name: str) -> None: - code = FakeCode("profiler", 0, name) - pframe = self.cur[-2] if self.cur else None - frame = FakeFrame(code, pframe) + class fake_code: + def __init__(self, filename, line, name): + self.co_filename = filename + self.co_line = line + self.co_name = name + self.co_firstlineno = 0 + + def __repr__(self): + return repr((self.co_filename, self.co_line, self.co_name, None)) + + class fake_frame: + def __init__(self, code, prior): + self.f_code = code + self.f_back = prior + self.f_locals = {} + + def simulate_call(self, name): + code = self.fake_code("profiler", 0, name) + if self.cur: + pframe = self.cur[-2] + else: + pframe = None + frame = self.fake_frame(code, pframe) self.dispatch["call"](self, frame, 0) - def simulate_cmd_complete(self) -> None: + def simulate_cmd_complete(self): get_time = self.timer t = get_time() - self.t while self.cur[-1]: @@ -511,174 +484,60 @@ def simulate_cmd_complete(self) -> None: t = 0 self.t = get_time() - t - def print_stats(self, sort: str | int | tuple = -1) -> None: - if not self.stats: - console.print("Codeflash: No stats available to print") - self.total_tt = 0 - return + def print_stats(self, sort=-1): + import pstats if not isinstance(sort, tuple): sort = (sort,) - - # First, convert stats to make them pstats-compatible - try: - # Initialize empty collections for pstats - self.files = [] - self.top_level = [] - - # Create entirely new dictionaries instead of modifying existing ones - new_stats = {} - new_timings = {} - - # Convert stats dictionary - stats_items = list(self.stats.items()) - for func, stats_data in stats_items: - try: - # Make sure we have 5 elements in stats_data - if len(stats_data) != 5: - console.print(f"Skipping malformed stats data for {func}: {stats_data}") - continue - - cc, nc, tt, ct, callers = stats_data - - if len(func) == 4: - file_name, line_num, func_name, class_name = func - new_func_name = f"{class_name}.{func_name}" if class_name else func_name - new_func = (file_name, line_num, new_func_name) - else: - new_func = func # Keep as is if already in correct format - - new_callers = {} - callers_items = list(callers.items()) - for caller_func, count in callers_items: - if isinstance(caller_func, tuple): - if len(caller_func) == 4: - caller_file, caller_line, caller_name, caller_class = caller_func - caller_new_name = f"{caller_class}.{caller_name}" if caller_class else caller_name - new_caller_func = (caller_file, caller_line, caller_new_name) - else: - new_caller_func = caller_func - else: - console.print(f"Unexpected caller format: {caller_func}") - new_caller_func = str(caller_func) - - new_callers[new_caller_func] = count - - # Store with new format - new_stats[new_func] = (cc, nc, tt, ct, new_callers) - except Exception as e: # noqa: BLE001 - console.print(f"Error converting stats for {func}: {e}") - continue - - timings_items = list(self.timings.items()) - for func, timing_data in timings_items: - try: - if len(timing_data) != 5: - console.print(f"Skipping malformed timing data for {func}: {timing_data}") - continue - - cc, ns, tt, ct, callers = timing_data - - if len(func) == 4: - file_name, line_num, func_name, class_name = func - new_func_name = f"{class_name}.{func_name}" if class_name else func_name - new_func = (file_name, line_num, new_func_name) - else: - new_func = func - - new_callers = {} - callers_items = list(callers.items()) - for caller_func, count in callers_items: - if isinstance(caller_func, tuple): - if len(caller_func) == 4: - caller_file, caller_line, caller_name, caller_class = caller_func - caller_new_name = f"{caller_class}.{caller_name}" if caller_class else caller_name - new_caller_func = (caller_file, caller_line, caller_new_name) - else: - new_caller_func = caller_func - else: - console.print(f"Unexpected caller format: {caller_func}") - new_caller_func = str(caller_func) - - new_callers[new_caller_func] = count - - new_timings[new_func] = (cc, ns, tt, ct, new_callers) - except Exception as e: # noqa: BLE001 - console.print(f"Error converting timings for {func}: {e}") - continue - - self.stats = new_stats - self.timings = new_timings - - self.total_tt = sum(tt for _, _, tt, _, _ in self.stats.values()) - - total_calls = sum(cc for cc, _, _, _, _ in self.stats.values()) - total_primitive = sum(nc for _, nc, _, _, _ in self.stats.values()) - - summary = Text.assemble( - f"{total_calls:,} function calls ", - ("(" + f"{total_primitive:,} primitive calls" + ")", "dim"), - f" in {self.total_tt / 1e6:.3f}milliseconds", - ) - - console.print(Align.center(Panel(summary, border_style="blue", width=80, padding=(0, 2), expand=False))) - - table = Table( - show_header=True, - header_style="bold magenta", - border_style="blue", - title="[bold]Function Profile[/bold] (ordered by internal time)", - title_style="cyan", - caption=f"Showing top 25 of {len(self.stats)} functions", - ) - - table.add_column("Calls", justify="right", style="green", width=10) - table.add_column("Time (ms)", justify="right", style="cyan", width=10) - table.add_column("Per Call", justify="right", style="cyan", width=10) - table.add_column("Cum (ms)", justify="right", style="yellow", width=10) - table.add_column("Cum/Call", justify="right", style="yellow", width=10) - table.add_column("Function", style="blue") - - sorted_stats = sorted( - ((func, stats) for func, stats in self.stats.items() if isinstance(func, tuple) and len(func) == 3), - key=lambda x: x[1][2], # Sort by tt (internal time) - reverse=True, - )[:25] # Limit to top 25 - - # Format and add each row to the table - for func, (cc, nc, tt, ct, _) in sorted_stats: - filename, lineno, funcname = func - - # Format calls - show recursive format if different - calls_str = f"{cc}/{nc}" if cc != nc else f"{cc:,}" - - # Convert to milliseconds - tt_ms = tt / 1e6 - ct_ms = ct / 1e6 - - # Calculate per-call times - per_call = tt_ms / cc if cc > 0 else 0 - cum_per_call = ct_ms / nc if nc > 0 else 0 - base_filename = Path(filename).name - file_link = f"[link=file://{filename}]{base_filename}[/link]" - - table.add_row( - calls_str, - f"{tt_ms:.3f}", - f"{per_call:.3f}", - f"{ct_ms:.3f}", - f"{cum_per_call:.3f}", - f"{funcname} [dim]({file_link}:{lineno})[/dim]", + # The following code customizes the default printing behavior to + # print in milliseconds. + s = StringIO() + stats_obj = pstats.Stats(copy(self), stream=s) + stats_obj.strip_dirs().sort_stats(*sort).print_stats(100) + self.total_tt = stats_obj.total_tt + console.print("total_tt", self.total_tt) + raw_stats = s.getvalue() + m = re.search(r"function calls?.*in (\d+)\.\d+ (seconds?)", raw_stats) + total_time = None + if m: + total_time = int(m.group(1)) + if total_time is None: + console.print("Failed to get total time from stats") + total_time_ms = total_time / 1e6 + raw_stats = re.sub( + r"(function calls?.*)in (\d+)\.\d+ (seconds?)", rf"\1 in {total_time_ms:.3f} milliseconds", raw_stats + ) + match_pattern = r"^ *[\d\/]+ +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +" + m = re.findall(match_pattern, raw_stats, re.MULTILINE) + ms_times = [] + for tottime, percall, cumtime, percall_cum in m: + tottime_ms = int(tottime) / 1e6 + percall_ms = int(percall) / 1e6 + cumtime_ms = int(cumtime) / 1e6 + percall_cum_ms = int(percall_cum) / 1e6 + ms_times.append([tottime_ms, percall_ms, cumtime_ms, percall_cum_ms]) + split_stats = raw_stats.split("\n") + new_stats = [] + + replace_pattern = r"^( *[\d\/]+) +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +(.*)" + times_index = 0 + for line in split_stats: + if times_index >= len(ms_times): + replaced = line + else: + replaced, n = re.subn( + replace_pattern, + rf"\g<1>{ms_times[times_index][0]:8.3f} {ms_times[times_index][1]:8.3f} {ms_times[times_index][2]:8.3f} {ms_times[times_index][3]:8.3f} \g<6>", + line, + count=1, ) + if n > 0: + times_index += 1 + new_stats.append(replaced) - console.print(Align.center(table)) - - except Exception as e: # noqa: BLE001 - console.print(f"[bold red]Error in stats processing:[/bold red] {e}") - console.print(f"Traced {self.trace_count:,} function calls") - self.total_tt = 0 + console.print("\n".join(new_stats)) - def make_pstats_compatible(self) -> None: + def make_pstats_compatible(self): # delete the extra class_name item from the function tuple self.files = [] self.top_level = [] @@ -693,33 +552,36 @@ def make_pstats_compatible(self) -> None: self.stats = new_stats self.timings = new_timings - def dump_stats(self, file: str) -> None: - with Path(file).open("wb") as f: + def dump_stats(self, file): + with open(file, "wb") as f: + self.create_stats() marshal.dump(self.stats, f) - def create_stats(self) -> None: + def create_stats(self): self.simulate_cmd_complete() self.snapshot_stats() - def snapshot_stats(self) -> None: + def snapshot_stats(self): self.stats = {} - for func, (cc, _ns, tt, ct, caller_dict) in self.timings.items(): - callers = caller_dict.copy() + for func, (cc, ns, tt, ct, callers) in self.timings.items(): + callers = callers.copy() nc = 0 for callcnt in callers.values(): nc += callcnt self.stats[func] = cc, nc, tt, ct, callers - def runctx(self, cmd: str, global_vars: dict[str, Any], local_vars: dict[str, Any]) -> Tracer | None: + def runctx(self, cmd, globals, locals): self.__enter__() try: - exec(cmd, global_vars, local_vars) # noqa: S102 + exec(cmd, globals, locals) finally: self.__exit__(None, None, None) return self -def main() -> ArgumentParser: +def main(): + from argparse import ArgumentParser + parser = ArgumentParser(allow_abbrev=False) parser.add_argument("-o", "--outfile", dest="outfile", help="Save trace to ", required=True) parser.add_argument("--only-functions", help="Trace only these functions", nargs="+", default=None) @@ -776,13 +638,16 @@ def main() -> ArgumentParser: "__cached__": None, } try: - Tracer( + tracer = Tracer( output=args.outfile, functions=args.only_functions, max_function_count=args.max_function_count, timeout=args.tracer_timeout, config_file_path=args.codeflash_config, - ).runctx(code, globs, None) + ) + + tracer.runctx(code, globs, None) + print(tracer.functions) except BrokenPipeError as exc: # Prevent "Exception ignored" during interpreter shutdown. diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py index 99151f983..db01ff049 100644 --- a/codeflash/verification/test_results.py +++ b/codeflash/verification/test_results.py @@ -29,7 +29,6 @@ class TestType(Enum): REPLAY_TEST = 4 CONCOLIC_COVERAGE_TEST = 5 INIT_STATE_TEST = 6 - BENCHMARK_TEST = 7 def to_name(self) -> str: if self == TestType.INIT_STATE_TEST: @@ -40,7 +39,6 @@ def to_name(self) -> str: TestType.GENERATED_REGRESSION: "πŸŒ€ Generated Regression Tests", TestType.REPLAY_TEST: "βͺ Replay Tests", TestType.CONCOLIC_COVERAGE_TEST: "πŸ”Ž Concolic Coverage Tests", - TestType.BENCHMARK_TEST: "πŸ“ Benchmark Tests", } return names[self] diff --git a/codeflash/verification/verification_utils.py b/codeflash/verification/verification_utils.py index 79f1b9656..99aaac939 100644 --- a/codeflash/verification/verification_utils.py +++ b/codeflash/verification/verification_utils.py @@ -75,3 +75,4 @@ class TestConfig: # or for unittest - project_root_from_module_root(args.tests_root, pyproject_file_path) concolic_test_root_dir: Optional[Path] = None pytest_cmd: str = "pytest" + benchmark_tests_root: Optional[Path] = None diff --git a/pyproject.toml b/pyproject.toml index a181fac2e..026c0eafd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,7 +69,7 @@ exclude = [ [tool.poetry.dependencies] python = ">=3.9" unidiff = ">=0.7.4" -pytest = ">=7.0.0,<8.3.4" +pytest = ">=7.0.0" gitpython = ">=3.1.31" libcst = ">=1.0.1" jedi = ">=0.19.1" @@ -92,7 +92,6 @@ rich = ">=13.8.1" lxml = ">=5.3.0" crosshair-tool = ">=0.0.78" coverage = ">=7.6.4" -line_profiler=">=4.2.0" #this is the minimum version which supports python 3.13 [tool.poetry.group.dev] optional = true @@ -120,7 +119,7 @@ types-gevent = "^24.11.0.20241230" types-greenlet = "^3.1.0.20241221" types-pexpect = "^4.9.0.20241208" types-unidiff = "^0.7.0.20240505" -uv = ">=0.6.2" +sqlalchemy = "^2.0.38" [tool.poetry.build] script = "codeflash/update_license_version.py" @@ -152,7 +151,7 @@ warn_required_dynamic_aliases = true line-length = 120 fix = true show-fixes = true -exclude = ["code_to_optimize/", "pie_test_set/", "tests/"] +exclude = ["code_to_optimize/", "pie_test_set/"] [tool.ruff.lint] select = ["ALL"] @@ -164,11 +163,10 @@ ignore = [ "D103", "D105", "D107", - "D203", # incorrect-blank-line-before-class (incompatible with D211) - "D213", # multi-line-summary-second-line (incompatible with D212) "S101", "S603", "S607", + "ANN101", "COM812", "FIX002", "PLR0912", @@ -177,14 +175,13 @@ ignore = [ "TD002", "TD003", "TD004", - "PLR2004", - "UP007" # remove once we drop 3.9 support. + "PLR2004" ] [tool.ruff.lint.flake8-type-checking] strict = true runtime-evaluated-base-classes = ["pydantic.BaseModel"] -runtime-evaluated-decorators = ["pydantic.validate_call", "pydantic.dataclasses.dataclass"] +runtime-evaluated-decorators = ["pydantic.validate_call"] [tool.ruff.lint.pep8-naming] classmethod-decorators = [ @@ -192,9 +189,6 @@ classmethod-decorators = [ "pydantic.validator", ] -[tool.ruff.lint.isort] -split-on-trailing-comma = false - [tool.ruff.format] docstring-code-format = true skip-magic-trailing-comma = true @@ -217,13 +211,13 @@ initial-content = """ [tool.codeflash] -module-root = "codeflash" -tests-root = "tests" +# All paths are relative to this pyproject.toml's directory. +module-root = "code_to_optimize" +tests-root = "code_to_optimize/tests" +benchmarks-root = "code_to_optimize/tests/pytest/benchmarks" test-framework = "pytest" -formatter-cmds = [ - "uvx ruff check --exit-zero --fix $file", - "uvx ruff format $file", -] +ignore-paths = [] +formatter-cmds = ["ruff check --exit-zero --fix $file", "ruff format $file"] [build-system] diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py new file mode 100644 index 000000000..acd40b0b3 --- /dev/null +++ b/tests/test_trace_benchmarks.py @@ -0,0 +1,8 @@ +from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest +from pathlib import Path + +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"]) \ No newline at end of file diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index 1aad04d42..4bf99c049 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -3,6 +3,7 @@ from pathlib import Path from codeflash.discovery.discover_unit_tests import discover_unit_tests +from codeflash.verification.test_results import TestType from codeflash.verification.verification_utils import TestConfig @@ -21,7 +22,7 @@ def test_unit_test_discovery_pytest(): def test_benchmark_test_discovery_pytest(): project_path = Path(__file__).parent.parent.resolve() / "code_to_optimize" - tests_path = project_path / "tests" / "pytest" + tests_path = project_path / "tests" / "pytest" / "benchmarks" / "test_benchmark_bubble_sort.py" test_config = TestConfig( tests_root=tests_path, project_root_path=project_path, @@ -29,9 +30,10 @@ def test_benchmark_test_discovery_pytest(): tests_project_rootdir=tests_path.parent, ) tests = discover_unit_tests(test_config) - print(tests) assert len(tests) > 0 - # print(tests) + assert 'bubble_sort.sorter' in tests + benchmark_tests = sum(1 for test in tests['bubble_sort.sorter'] if test.tests_in_file.test_type == TestType.BENCHMARK_TEST) + assert benchmark_tests == 1 def test_unit_test_discovery_unittest(): From 7590c29dff17e01edb10d518c018d5ee0493eb2d Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 4 Mar 2025 15:53:24 -0800 Subject: [PATCH 056/122] initial implementation of tracing benchmarks via the plugin --- codeflash/benchmarking/get_trace_info.py | 10 ++-- codeflash/benchmarking/plugin/plugin.py | 3 +- .../pytest_new_process_trace_benchmarks.py | 2 +- codeflash/benchmarking/utils.py | 26 ++++++++++ codeflash/cli_cmds/cli.py | 1 + codeflash/code_utils/config_parser.py | 4 +- codeflash/models/models.py | 4 +- codeflash/optimization/function_optimizer.py | 51 +++++++++++++++---- codeflash/optimization/optimizer.py | 13 +++-- codeflash/tracer.py | 40 ++++++++++----- codeflash/verification/test_results.py | 29 +++++++++++ codeflash/verification/test_runner.py | 2 + 12 files changed, 149 insertions(+), 36 deletions(-) create mode 100644 codeflash/benchmarking/utils.py diff --git a/codeflash/benchmarking/get_trace_info.py b/codeflash/benchmarking/get_trace_info.py index 1d0b339d9..3dd3831ce 100644 --- a/codeflash/benchmarking/get_trace_info.py +++ b/codeflash/benchmarking/get_trace_info.py @@ -54,18 +54,20 @@ def get_function_benchmark_timings(trace_dir: Path, all_functions_to_optimize: l # Adjust query based on whether we have a class name if class_name: cursor.execute( - "SELECT total_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND class_name = ?", + "SELECT cumulative_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND class_name = ?", (f"%{filename}", function_name, class_name) ) else: cursor.execute( - "SELECT total_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND (class_name IS NULL OR class_name = '')", + "SELECT cumulative_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND (class_name IS NULL OR class_name = '')", (f"%{filename}", function_name) ) - result = cursor.fetchone() + result = cursor.fetchall() + if len(result) > 1: + print(f"Multiple results found for {qualified_name} in {benchmark_name}: {result}") if result: - time_ns = result[0] + time_ns = result[0][0] function_benchmark_timings[qualified_name][benchmark_name] = time_ns / 1e6 # Convert to milliseconds conn.close() diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index 34ca2b777..80accec22 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -68,7 +68,8 @@ def __call__(self, func, *args, **kwargs): tracer = Tracer( output=str(output_path), # Convert Path to string for Tracer functions=trace_functions, - max_function_count=256 + max_function_count=256, + benchmark=True ) with tracer: diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index b892d62a0..6b91e2b4f 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -8,7 +8,7 @@ try: exitcode = pytest.main( - [benchmarks_root, "--benchmarks-root", benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "--functions", function_list], plugins=[CodeFlashPlugin()] + [benchmarks_root, "--benchmarks-root", benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "--functions", function_list,"-o", "addopts="], plugins=[CodeFlashPlugin()] ) except Exception as e: print(f"Failed to collect tests: {e!s}") diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py new file mode 100644 index 000000000..d97c2e36e --- /dev/null +++ b/codeflash/benchmarking/utils.py @@ -0,0 +1,26 @@ +def print_benchmark_table(function_benchmark_timings, total_benchmark_timings): + # Print table header + print(f"{'Benchmark Test':<50} | {'Total Time (s)':<15} | {'Function Time (s)':<15} | {'Percentage (%)':<15}") + print("-" * 100) + + # Process each function's benchmark data + for func_path, test_times in function_benchmark_timings.items(): + function_name = func_path.split(":")[-1] + print(f"\n== Function: {function_name} ==") + + # Sort by percentage (highest first) + sorted_tests = [] + for test_name, func_time in test_times.items(): + total_time = total_benchmark_timings.get(test_name, 0) + if total_time > 0: + percentage = (func_time / total_time) * 100 + sorted_tests.append((test_name, total_time, func_time, percentage)) + + sorted_tests.sort(key=lambda x: x[3], reverse=True) + + # Print each test's data + for test_name, total_time, func_time, percentage in sorted_tests: + print(f"{test_name:<50} | {total_time:<15.3f} | {func_time:<15.3f} | {percentage:<15.2f}") + +# Usage + diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 04445f1db..96bb0cef3 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -113,6 +113,7 @@ def process_pyproject_config(args: Namespace) -> Namespace: supported_keys = [ "module_root", "tests_root", + "benchmarks_root", "test_framework", "ignore_paths", "pytest_cmd", diff --git a/codeflash/code_utils/config_parser.py b/codeflash/code_utils/config_parser.py index d814f12d0..6f2b1268c 100644 --- a/codeflash/code_utils/config_parser.py +++ b/codeflash/code_utils/config_parser.py @@ -52,10 +52,10 @@ def parse_config_file( assert isinstance(config, dict) # default values: - path_keys = ["module-root", "tests-root"] + path_keys = ["module-root", "tests-root", "benchmarks-root"] path_list_keys = ["ignore-paths"] str_keys = {"pytest-cmd": "pytest", "git-remote": "origin"} - bool_keys = {"disable-telemetry": False, "disable-imports-sorting": False} + bool_keys = {"disable-telemetry": False, "disable-imports-sorting": False, "benchmark": False} list_str_keys = {"formatter-cmds": ["black $file"]} for key in str_keys: diff --git a/codeflash/models/models.py b/codeflash/models/models.py index a00834cdd..c0ce74c47 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -76,8 +76,10 @@ class BestOptimization(BaseModel): candidate: OptimizedCandidate helper_functions: list[FunctionSource] runtime: int + replay_runtime: int | None winning_behavioral_test_results: TestResults winning_benchmarking_test_results: TestResults + winning_replay_benchmarking_test_results : TestResults | None = None class CodeString(BaseModel): @@ -224,7 +226,7 @@ class OriginalCodeBaseline(BaseModel): benchmarking_test_results: TestResults line_profile_results: dict runtime: int - coverage_results: Optional[CoverageData] + coverage_results: CoverageData | None class CoverageStatus(Enum): diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 66d3c6ab6..38277851b 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -117,7 +117,6 @@ def __init__( self.experiment_id = os.getenv("CODEFLASH_EXPERIMENT_ID", None) self.local_aiservice_client = LocalAiServiceClient() if self.experiment_id else None self.test_files = TestFiles(test_files=[]) - self.args = args # Check defaults for these self.function_trace_id: str = str(uuid.uuid4()) self.original_module_path = module_name_from_file_path(self.function_to_optimize.file_path, self.project_root) @@ -285,20 +284,22 @@ def optimize_function(self) -> Result[BestOptimization, str]: function_name=function_to_optimize_qualified_name, file_path=self.function_to_optimize.file_path, ) - speedup = explanation.speedup # eg. 1.2 means 1.2x faster + speedup = explanation.speedup # if self.args.benchmark: + original_replay_timing = original_code_baseline.benchmarking_test_results.total_replay_test_runtime() fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_file_name] for benchmark_name, og_benchmark_timing in fto_benchmark_timings.items(): print(f"Calculating speedup for benchmark {benchmark_name}") total_benchmark_timing = self.total_benchmark_timings[benchmark_name] # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values - expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + og_benchmark_timing / speedup + replay_speedup = original_replay_timing / best_optimization.replay_runtime - 1 + print(f"Replay speedup: {replay_speedup}") + expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / (replay_speedup + 1) * og_benchmark_timing print(f"Expected new benchmark timing: {expected_new_benchmark_timing}") print(f"Original benchmark timing: {total_benchmark_timing}") - print(f"Benchmark speedup: {total_benchmark_timing / expected_new_benchmark_timing}") - - speedup = total_benchmark_timing / expected_new_benchmark_timing - print(f"Speedup: {speedup}") + benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing + benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 + print(f"Benchmark speedup: {benchmark_speedup_percent:.2f}%") self.log_successful_optimization(explanation, generated_tests) @@ -447,13 +448,30 @@ def determine_best_candidate( ) tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") - + if self.args.benchmark: + original_code_replay_runtime = original_code_baseline.benchmarking_test_results.total_replay_test_runtime() + candidate_replay_runtime = candidate_result.benchmarking_test_results.total_replay_test_runtime() + replay_perf_gain = performance_gain( + original_runtime_ns=original_code_replay_runtime, + optimized_runtime_ns=candidate_replay_runtime, + ) + tree.add("Replay Benchmarking： ") + tree.add(f"Original summed runtime: {humanize_runtime(original_code_replay_runtime)}") + tree.add( + f"Best summed runtime: {humanize_runtime(candidate_replay_runtime)} " + f"(measured over {candidate_result.max_loop_count} " + f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + ) + tree.add(f"Speedup percentage: {replay_perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio: {replay_perf_gain + 1:.1f}X") best_optimization = BestOptimization( candidate=candidate, helper_functions=code_context.helper_functions, runtime=best_test_runtime, + replay_runtime=candidate_replay_runtime if self.args.benchmark else None, winning_behavioral_test_results=candidate_result.behavior_test_results, winning_benchmarking_test_results=candidate_result.benchmarking_test_results, + winning_replay_benchmarking_test_results=candidate_result.benchmarking_test_results, ) best_runtime_until_now = best_test_runtime else: @@ -664,6 +682,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi existing_test_files_count += 1 elif test_type == TestType.REPLAY_TEST: replay_test_files_count += 1 + print("Replay test found") elif test_type == TestType.CONCOLIC_COVERAGE_TEST: concolic_coverage_test_files_count += 1 else: @@ -708,6 +727,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi unique_instrumented_test_files.add(new_behavioral_test_path) unique_instrumented_test_files.add(new_perf_test_path) + if not self.test_files.get_by_original_file_path(path_obj_test_file): self.test_files.add( TestFile( @@ -719,6 +739,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi tests_in_file=[t.tests_in_file for t in tests_in_file_list], ) ) + logger.info( f"Discovered {existing_test_files_count} existing unit test file" f"{'s' if existing_test_files_count != 1 else ''}, {replay_test_files_count} replay test file" @@ -888,7 +909,6 @@ def establish_original_code_baseline( enable_coverage=False, code_context=code_context, ) - else: benchmarking_results = TestResults() start_time: float = time.time() @@ -917,7 +937,6 @@ def establish_original_code_baseline( ) console.rule() - total_timing = benchmarking_results.total_passed_runtime() # caution: doesn't handle the loop index functions_to_remove = [ result.id.test_function_name @@ -944,6 +963,9 @@ def establish_original_code_baseline( ) console.rule() logger.debug(f"Total original code runtime (ns): {total_timing}") + + if self.args.benchmark: + logger.info(f"Total replay test runtime: {humanize_runtime(benchmarking_results.total_replay_test_runtime())}") return Success( ( OriginalCodeBaseline( @@ -1062,6 +1084,15 @@ def run_optimized_candidate( console.rule() logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") + if self.args.benchmark: + total_candidate_replay_timing = ( + candidate_benchmarking_results.total_replay_test_runtime() + if candidate_benchmarking_results + else 0 + ) + logger.debug( + f"Total optimized code {optimization_candidate_index} replay benchmark runtime (ns): {total_candidate_replay_timing}" + ) return Success( OptimizedCandidateResult( max_loop_count=loop_count, diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 01a196143..9c5bc08ce 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -9,6 +9,7 @@ from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest +from codeflash.benchmarking.utils import print_benchmark_table from codeflash.cli_cmds.console import console, logger from codeflash.code_utils import env_utils from codeflash.code_utils.code_replacer import normalize_code, normalize_node @@ -23,7 +24,7 @@ from codeflash.verification.test_results import TestType 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 if TYPE_CHECKING: from argparse import Namespace @@ -98,11 +99,8 @@ def run(self) -> None: 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) - print(function_benchmark_timings) total_benchmark_timings = get_benchmark_timings(trace_dir) - print("Total benchmark timings:") - print(total_benchmark_timings) - # for function in fully_qualified_function_names: + print_benchmark_table(function_benchmark_timings, total_benchmark_timings) optimizations_found: int = 0 @@ -127,6 +125,7 @@ def run(self) -> None: console.rule() ph("cli-optimize-discovered-tests", {"num_tests": num_discovered_tests}) + for original_module_path in file_to_funcs_to_optimize: logger.info(f"Examining file {original_module_path!s}…") console.rule() @@ -217,6 +216,10 @@ def run(self) -> None: test_file.instrumented_behavior_file_path.unlink(missing_ok=True) if function_optimizer.test_cfg.concolic_test_root_dir: shutil.rmtree(function_optimizer.test_cfg.concolic_test_root_dir, ignore_errors=True) + if self.args.benchmark: + trace_dir = Path(self.args.benchmarks_root) / "codeflash_replay_tests" + if trace_dir.exists(): + shutil.rmtree(trace_dir, ignore_errors=True) if hasattr(get_run_tmp_file, "tmpdir"): get_run_tmp_file.tmpdir.cleanup() diff --git a/codeflash/tracer.py b/codeflash/tracer.py index 5bc1ae482..02a0e4157 100644 --- a/codeflash/tracer.py +++ b/codeflash/tracer.py @@ -58,6 +58,7 @@ def __init__( config_file_path: Path | None = None, max_function_count: int = 256, timeout: int | None = None, # seconds + benchmark: bool = False, ) -> None: """:param output: The path to the output trace file :param functions: List of functions to trace. If None, trace all functions @@ -95,7 +96,6 @@ def __init__( self.max_function_count = max_function_count self.config, found_config_path = parse_config_file(config_file_path) self.project_root = project_root_from_module_root(Path(self.config["module_root"]), found_config_path) - print("project_root", self.project_root) self.ignored_functions = {"", "", "", "", "", ""} self.file_being_called_from: str = str(Path(sys._getframe().f_back.f_code.co_filename).name).replace(".", "_") @@ -105,6 +105,7 @@ def __init__( self.next_insert = 1000 self.trace_count = 0 + self.benchmark = benchmark # Profiler variables self.bias = 0 # calibration constant self.timings = {} @@ -184,18 +185,25 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: cur.execute("INSERT INTO total_time VALUES (?)", (self.total_tt,)) self.con.commit() self.con.close() + function_string = [str(function.file_name) + ":" + (function.class_name + ":" if function.class_name else "") + function.function_name for function in self.function_modules] + # print(function_string) # filter any functions where we did not capture the return + # self.function_modules = [ + # function + # for function in self.function_modules + # if self.function_count[ + # str(function.file_name) + # + ":" + # + (function.class_name + ":" if function.class_name else "") + # + function.function_name + # ] + # > 0 + # ] self.function_modules = [ function for function in self.function_modules - if self.function_count[ - str(function.file_name) - + ":" - + (function.class_name + ":" if function.class_name else "") - + function.function_name - ] - > 0 + if str(str(function.file_name) + ":" + (function.class_name + ":" if function.class_name else "") + function.function_name) in self.function_count ] replay_test = create_trace_replay_test( @@ -207,13 +215,21 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: # Need a better way to store the replay test # function_path = "_".join(self.functions) if self.functions else self.file_being_called_from function_path = self.file_being_called_from - test_file_path = get_test_file_path( - test_dir=Path(self.config["tests_root"]), function_name=function_path, test_type="replay" - ) + if self.benchmark and self.config["benchmarks_root"]: + # check if replay test dir exists, create + replay_test_dir = Path(self.config["benchmarks_root"]) / "codeflash_replay_tests" + if not replay_test_dir.exists(): + replay_test_dir.mkdir(parents=True) + test_file_path = get_test_file_path( + test_dir=replay_test_dir, function_name=function_path, test_type="replay" + ) + else: + test_file_path = get_test_file_path( + test_dir=Path(self.config["tests_root"]), function_name=function_path, test_type="replay" + ) replay_test = isort.code(replay_test) with open(test_file_path, "w", encoding="utf8") as file: file.write(replay_test) - console.print( f"Codeflash: Traced {self.trace_count} function calls successfully and replay test created at - {test_file_path}", crop=False, diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py index db01ff049..916f6da11 100644 --- a/codeflash/verification/test_results.py +++ b/codeflash/verification/test_results.py @@ -193,6 +193,35 @@ def total_passed_runtime(self) -> int: ] ) + def usable_replay_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: + """Collect runtime data for replay tests that passed and have runtime information. + + :return: A dictionary mapping invocation IDs to lists of runtime values. + """ + usable_runtimes = [ + (result.id, result.runtime) + for result in self.test_results + if result.did_pass and result.runtime and result.test_type == TestType.REPLAY_TEST + ] + + return { + usable_id: [runtime[1] for runtime in usable_runtimes if runtime[0] == usable_id] + for usable_id in {runtime[0] for runtime in usable_runtimes} + } + + def total_replay_test_runtime(self) -> int: + """Calculate the sum of runtimes of replay test cases that passed, where a testcase runtime + is the minimum value of all looped execution runtimes. + + :return: The runtime in nanoseconds. + """ + replay_runtime_data = self.usable_replay_runtime_data_by_test_case() + + return sum([ + min(runtimes) + for invocation_id, runtimes in replay_runtime_data.items() + ]) if replay_runtime_data else 0 + def __iter__(self) -> Iterator[FunctionTestInvocation]: return iter(self.test_results) diff --git a/codeflash/verification/test_runner.py b/codeflash/verification/test_runner.py index d8c58eb9a..852b6bf8a 100644 --- a/codeflash/verification/test_runner.py +++ b/codeflash/verification/test_runner.py @@ -62,6 +62,8 @@ def run_behavioral_tests( "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", + "-o", + "addopts=", "--codeflash_loops_scope=session", "--codeflash_min_loops=1", "--codeflash_max_loops=1", From 034bed3075d961e3691214d4635279e9e941632b Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 11 Mar 2025 14:37:44 -0700 Subject: [PATCH 057/122] basic version working on bubble sort --- code_to_optimize/bubble_sort.py | 12 ++----- .../benchmarks/test_process_and_sort.py | 2 +- .../tests/unittest/test_bubble_sort.py | 36 +++++++++---------- .../unittest/test_bubble_sort_parametrized.py | 36 +++++++++---------- codeflash/optimization/optimizer.py | 2 ++ pyproject.toml | 22 +++++++----- 6 files changed, 56 insertions(+), 54 deletions(-) diff --git a/code_to_optimize/bubble_sort.py b/code_to_optimize/bubble_sort.py index 787cc4a90..fd53c04a7 100644 --- a/code_to_optimize/bubble_sort.py +++ b/code_to_optimize/bubble_sort.py @@ -1,10 +1,4 @@ def sorter(arr): - print("codeflash stdout: Sorting list") - for i in range(len(arr)): - for j in range(len(arr) - 1): - if arr[j] > arr[j + 1]: - temp = arr[j] - arr[j] = arr[j + 1] - arr[j + 1] = temp - print(f"result: {arr}") - return arr \ No newline at end of file + # Utilizing Python's built-in Timsort algorithm for better performance + arr.sort() + return arr diff --git a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py index ca2f0ef65..93d78afef 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py @@ -1,5 +1,5 @@ from code_to_optimize.process_and_bubble_sort import compute_and_sort -from code_to_optimize.bubble_sort2 import sorter +from code_to_optimize.bubble_sort import sorter def test_compute_and_sort(benchmark): result = benchmark(compute_and_sort, list(reversed(range(5000)))) assert result == 6247083.5 diff --git a/code_to_optimize/tests/unittest/test_bubble_sort.py b/code_to_optimize/tests/unittest/test_bubble_sort.py index 200f82b7a..4c76414ef 100644 --- a/code_to_optimize/tests/unittest/test_bubble_sort.py +++ b/code_to_optimize/tests/unittest/test_bubble_sort.py @@ -1,18 +1,18 @@ -import unittest - -from code_to_optimize.bubble_sort import sorter - - -class TestPigLatin(unittest.TestCase): - def test_sort(self): - input = [5, 4, 3, 2, 1, 0] - output = sorter(input) - self.assertEqual(output, [0, 1, 2, 3, 4, 5]) - - input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0] - output = sorter(input) - self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) - - input = list(reversed(range(5000))) - output = sorter(input) - self.assertEqual(output, list(range(5000))) +# import unittest +# +# from code_to_optimize.bubble_sort import sorter +# +# +# class TestPigLatin(unittest.TestCase): +# def test_sort(self): +# input = [5, 4, 3, 2, 1, 0] +# output = sorter(input) +# self.assertEqual(output, [0, 1, 2, 3, 4, 5]) +# +# input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0] +# output = sorter(input) +# self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) +# +# input = list(reversed(range(5000))) +# output = sorter(input) +# self.assertEqual(output, list(range(5000))) diff --git a/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py b/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py index 59c86abc8..c1aef993b 100644 --- a/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py +++ b/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py @@ -1,18 +1,18 @@ -import unittest - -from parameterized import parameterized - -from code_to_optimize.bubble_sort import sorter - - -class TestPigLatin(unittest.TestCase): - @parameterized.expand( - [ - ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]), - ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), - (list(reversed(range(50))), list(range(50))), - ] - ) - def test_sort(self, input, expected_output): - output = sorter(input) - self.assertEqual(output, expected_output) +# import unittest +# +# from parameterized import parameterized +# +# from code_to_optimize.bubble_sort import sorter +# +# +# class TestPigLatin(unittest.TestCase): +# @parameterized.expand( +# [ +# ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]), +# ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), +# (list(reversed(range(50))), list(range(50))), +# ] +# ) +# def test_sort(self, input, expected_output): +# output = sorter(input) +# self.assertEqual(output, expected_output) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 9c5bc08ce..11d44a349 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -78,6 +78,8 @@ def run(self) -> None: function_optimizer = None file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]] num_optimizable_functions: int + # if self.args.benchmark: + # discover functions (file_to_funcs_to_optimize, num_optimizable_functions) = get_functions_to_optimize( optimize_all=self.args.all, replay_test=self.args.replay_test, diff --git a/pyproject.toml b/pyproject.toml index 026c0eafd..877815004 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,7 +69,7 @@ exclude = [ [tool.poetry.dependencies] python = ">=3.9" unidiff = ">=0.7.4" -pytest = ">=7.0.0" +pytest = ">=7.0.0,<8.3.4" gitpython = ">=3.1.31" libcst = ">=1.0.1" jedi = ">=0.19.1" @@ -119,7 +119,7 @@ types-gevent = "^24.11.0.20241230" types-greenlet = "^3.1.0.20241221" types-pexpect = "^4.9.0.20241208" types-unidiff = "^0.7.0.20240505" -sqlalchemy = "^2.0.38" +uv = ">=0.6.2" [tool.poetry.build] script = "codeflash/update_license_version.py" @@ -163,10 +163,11 @@ ignore = [ "D103", "D105", "D107", + "D203", # incorrect-blank-line-before-class (incompatible with D211) + "D213", # multi-line-summary-second-line (incompatible with D212) "S101", "S603", "S607", - "ANN101", "COM812", "FIX002", "PLR0912", @@ -175,13 +176,14 @@ ignore = [ "TD002", "TD003", "TD004", - "PLR2004" + "PLR2004", + "UP007" # remove once we drop 3.9 support. ] [tool.ruff.lint.flake8-type-checking] strict = true runtime-evaluated-base-classes = ["pydantic.BaseModel"] -runtime-evaluated-decorators = ["pydantic.validate_call"] +runtime-evaluated-decorators = ["pydantic.validate_call", "pydantic.dataclasses.dataclass"] [tool.ruff.lint.pep8-naming] classmethod-decorators = [ @@ -189,6 +191,9 @@ classmethod-decorators = [ "pydantic.validator", ] +[tool.ruff.lint.isort] +split-on-trailing-comma = false + [tool.ruff.format] docstring-code-format = true skip-magic-trailing-comma = true @@ -211,13 +216,14 @@ initial-content = """ [tool.codeflash] -# All paths are relative to this pyproject.toml's directory. module-root = "code_to_optimize" tests-root = "code_to_optimize/tests" benchmarks-root = "code_to_optimize/tests/pytest/benchmarks" test-framework = "pytest" -ignore-paths = [] -formatter-cmds = ["ruff check --exit-zero --fix $file", "ruff format $file"] +formatter-cmds = [ + "uvx ruff check --exit-zero --fix $file", + "uvx ruff format $file", +] [build-system] From 1f3fd4d2dba1b6dc8ada99e4fb34f2016ac6596d Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 11 Mar 2025 16:00:21 -0700 Subject: [PATCH 058/122] initial attempt for codeflash_trace_decorator --- .../benchmarking/codeflash_trace_decorator.py | 73 +++++++++++++++++++ .../instrument_codeflash_trace.py | 1 + codeflash/benchmarking/plugin/plugin.py | 46 ++---------- .../pytest_new_process_trace_benchmarks.py | 3 +- codeflash/benchmarking/trace_benchmarks.py | 5 +- codeflash/optimization/optimizer.py | 35 ++++++--- tests/test_trace_benchmarks.py | 2 +- 7 files changed, 106 insertions(+), 59 deletions(-) create mode 100644 codeflash/benchmarking/codeflash_trace_decorator.py create mode 100644 codeflash/benchmarking/instrument_codeflash_trace.py diff --git a/codeflash/benchmarking/codeflash_trace_decorator.py b/codeflash/benchmarking/codeflash_trace_decorator.py new file mode 100644 index 000000000..f996fa295 --- /dev/null +++ b/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 diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -0,0 +1 @@ + diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index 80accec22..bb903e554 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/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() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 6b91e2b4f..85a6755bf 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/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}") diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 2d3acdd66..bec5a03d4 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/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, diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 11d44a349..8d28e8c9d 100644 --- a/codeflash/optimization/optimizer.py +++ b/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 diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index acd40b0b3..688e23d8c 100644 --- a/tests/test_trace_benchmarks.py +++ b/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"]) \ No newline at end of file + trace_benchmarks_pytest(benchmarks_root, project_root) \ No newline at end of file From 5faccd821c8e21677f3b91d3c805cfca27b76dd9 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 11 Mar 2025 17:35:34 -0700 Subject: [PATCH 059/122] improvements --- code_to_optimize/bubble_sort.py | 10 +- codeflash/benchmarking/codeflash_trace.py | 122 ++++++++++++++++++ .../benchmarking/codeflash_trace_decorator.py | 73 ----------- codeflash/benchmarking/plugin/plugin.py | 8 +- tests/test_codeflash_trace_decorator.py | 15 +++ 5 files changed, 150 insertions(+), 78 deletions(-) create mode 100644 codeflash/benchmarking/codeflash_trace.py delete mode 100644 codeflash/benchmarking/codeflash_trace_decorator.py create mode 100644 tests/test_codeflash_trace_decorator.py diff --git a/code_to_optimize/bubble_sort.py b/code_to_optimize/bubble_sort.py index fd53c04a7..41cca9cea 100644 --- a/code_to_optimize/bubble_sort.py +++ b/code_to_optimize/bubble_sort.py @@ -1,4 +1,10 @@ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +@codeflash_trace("bubble_sort.trace") def sorter(arr): - # Utilizing Python's built-in Timsort algorithm for better performance - arr.sort() + for i in range(len(arr)): + for j in range(len(arr) - 1): + if arr[j] > arr[j + 1]: + temp = arr[j] + arr[j] = arr[j + 1] + arr[j + 1] = temp return arr diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py new file mode 100644 index 000000000..428f4a22c --- /dev/null +++ b/codeflash/benchmarking/codeflash_trace.py @@ -0,0 +1,122 @@ +import functools +import os +import pickle +import sqlite3 +import time +from typing import Callable + + +class CodeflashTrace: + """A class that provides both a decorator for tracing function calls + and a context manager for managing the tracing data lifecycle. + """ + + def __init__(self) -> None: + self.function_calls_data = [] + + def __enter__(self) -> None: + # Initialize for context manager use + self.function_calls_data = [] + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + # Cleanup is optional here + pass + + def __call__(self, func: Callable) -> Callable: + """Use as a decorator to trace function execution. + + Args: + func: The function to be decorated + + Returns: + The wrapped function + + """ + @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() + overhead_time = 0 + + try: + # 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", "") + + # Calculate overhead time + overhead_end_time = time.time() + overhead_time = overhead_end_time - overhead_start_time + + self.function_calls_data.append( + (func.__name__, func.__module__, func.__code__.co_filename, + benchmark_function_name, benchmark_file_name, execution_time, + overhead_time, pickled_args, pickled_kwargs) + ) + + except Exception as e: + print(f"Error in codeflash_trace: {e}") + + return result + return wrapper + + def write_to_db(self, output_file: str) -> None: + """Write all collected function call data to the SQLite database. + + Args: + output_file: Path to the SQLite database file where results will be stored + + """ + if not self.function_calls_data: + print("No function call data to write") + return + + 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, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" + ) + + # Insert all data at once + cur.executemany( + "INSERT INTO function_calls " + "(function_name, class_name, file_name, benchmark_function_name, " + "benchmark_file_name, time_ns, overhead_time_ns, args, kwargs) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + self.function_calls_data + ) + + # Commit and close + con.commit() + con.close() + + print(f"Successfully wrote {len(self.function_calls_data)} function call records to {output_file}") + + # Clear the data after writing + self.function_calls_data.clear() + + except Exception as e: + print(f"Error writing function calls to database: {e}") + +# Create a singleton instance +codeflash_trace = CodeflashTrace() diff --git a/codeflash/benchmarking/codeflash_trace_decorator.py b/codeflash/benchmarking/codeflash_trace_decorator.py deleted file mode 100644 index f996fa295..000000000 --- a/codeflash/benchmarking/codeflash_trace_decorator.py +++ /dev/null @@ -1,73 +0,0 @@ -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 diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index bb903e554..6d8db9bf9 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -1,6 +1,6 @@ import pytest import time - +import os class CodeFlashPlugin: @staticmethod def pytest_addoption(parser): @@ -35,9 +35,11 @@ def benchmark(request): class Benchmark: def __call__(self, func, *args, **kwargs): - start = time.time_ns() + os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = request.node.name + os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = request.node.fspath.basename + start = time.process_time_ns() result = func(*args, **kwargs) - end = time.time_ns() + end = time.process_time_ns() print(f"Benchmark: {func.__name__} took {end - start} ns") return result diff --git a/tests/test_codeflash_trace_decorator.py b/tests/test_codeflash_trace_decorator.py new file mode 100644 index 000000000..251b668ec --- /dev/null +++ b/tests/test_codeflash_trace_decorator.py @@ -0,0 +1,15 @@ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +from pathlib import Path + +@codeflash_trace("test_codeflash_trace.trace") +def example_function(arr): + arr.sort() + return arr + + +def test_codeflash_trace_decorator(): + arr = [3, 1, 2] + result = example_function(arr) + # cleanup test trace file using Path + assert result == [1, 2, 3] + Path("test_codeflash_trace.trace").unlink() \ No newline at end of file From d6217e8f8ec1a7867a0b3e80542aa7248ed3784f Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 12 Mar 2025 11:46:29 -0700 Subject: [PATCH 060/122] work on new replay_test logic --- code_to_optimize/bubble_sort.py | 2 +- codeflash/benchmarking/codeflash_trace.py | 70 +++++++- .../pytest_new_process_trace_benchmarks.py | 5 + codeflash/benchmarking/replay_test.py | 159 ++++++++++++++++++ codeflash/benchmarking/trace_benchmarks.py | 3 +- codeflash/optimization/optimizer.py | 7 + tests/test_codeflash_trace_decorator.py | 4 +- tests/test_trace_benchmarks.py | 6 +- 8 files changed, 246 insertions(+), 10 deletions(-) create mode 100644 codeflash/benchmarking/replay_test.py diff --git a/code_to_optimize/bubble_sort.py b/code_to_optimize/bubble_sort.py index 41cca9cea..91b77f50c 100644 --- a/code_to_optimize/bubble_sort.py +++ b/code_to_optimize/bubble_sort.py @@ -1,5 +1,5 @@ from codeflash.benchmarking.codeflash_trace import codeflash_trace -@codeflash_trace("bubble_sort.trace") +@codeflash_trace def sorter(arr): for i in range(len(arr)): for j in range(len(arr) - 1): diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 428f4a22c..45c7fa6c2 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -14,10 +14,10 @@ class CodeflashTrace: def __init__(self) -> None: self.function_calls_data = [] - def __enter__(self) -> None: - # Initialize for context manager use - self.function_calls_data = [] - return self + # def __enter__(self) -> None: + # # Initialize for context manager use + # self.function_calls_data = [] + # return self def __exit__(self, exc_type, exc_val, exc_tb) -> None: # Cleanup is optional here @@ -82,7 +82,7 @@ def write_to_db(self, output_file: str) -> None: if not self.function_calls_data: print("No function call data to write") return - + self.db_path = output_file try: # Connect to the database con = sqlite3.connect(output_file) @@ -118,5 +118,65 @@ def write_to_db(self, output_file: str) -> None: except Exception as e: print(f"Error writing function calls to database: {e}") + def print_codeflash_db(self, limit: int = None) -> None: + """ + Print the contents of a CodeflashTrace SQLite database. + + Args: + db_path: Path to the SQLite database file + limit: Maximum number of records to print (None for all) + """ + try: + # Connect to the database + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + + # Get the count of records + cursor.execute("SELECT COUNT(*) FROM function_calls") + total_records = cursor.fetchone()[0] + print(f"Found {total_records} function call records in {self.db_path}") + + # Build the query with optional limit + query = "SELECT * FROM function_calls" + if limit: + query += f" LIMIT {limit}" + + # Execute the query + cursor.execute(query) + + # Print column names + columns = [desc[0] for desc in cursor.description] + print("\nColumns:", columns) + print("\n" + "=" * 80 + "\n") + + # Print each row + for i, row in enumerate(cursor.fetchall()): + print(f"Record #{i + 1}:") + print(f" Function: {row[0]}") + print(f" Module: {row[1]}") + print(f" File: {row[2]}") + print(f" Benchmark Function: {row[3] or 'N/A'}") + print(f" Benchmark File: {row[4] or 'N/A'}") + print(f" Execution Time: {row[5]:.6f} seconds") + print(f" Overhead Time: {row[6]:.6f} seconds") + + # Unpickle and print args and kwargs + try: + args = pickle.loads(row[7]) + kwargs = pickle.loads(row[8]) + + print(f" Args: {args}") + print(f" Kwargs: {kwargs}") + except Exception as e: + print(f" Error unpickling args/kwargs: {e}") + print(f" Raw args: {row[7]}") + print(f" Raw kwargs: {row[8]}") + + print("\n" + "-" * 40 + "\n") + + conn.close() + + except Exception as e: + print(f"Error reading database: {e}") # Create a singleton instance codeflash_trace = CodeflashTrace() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 85a6755bf..ebe1fa4ae 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -1,7 +1,10 @@ import sys from plugin.plugin import CodeFlashPlugin +from codeflash.benchmarking.codeflash_trace import codeflash_trace +from codeflash.code_utils.code_utils import get_run_tmp_file benchmarks_root = sys.argv[1] +output_file = sys.argv[2] if __name__ == "__main__": import pytest @@ -9,6 +12,8 @@ exitcode = pytest.main( [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashPlugin()] ) + codeflash_trace.write_to_db(output_file) + codeflash_trace.print_codeflash_db() except Exception as e: print(f"Failed to collect tests: {e!s}") exitcode = -1 \ No newline at end of file diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py new file mode 100644 index 000000000..0bc2de1d4 --- /dev/null +++ b/codeflash/benchmarking/replay_test.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +import sqlite3 +import textwrap +from collections.abc import Generator +from typing import Any, List, Optional + +from codeflash.discovery.functions_to_optimize import FunctionProperties, inspect_top_level_functions_or_methods +from codeflash.tracing.tracing_utils import FunctionModules + + +def get_next_arg_and_return( + trace_file: str, function_name: str, file_name: str, class_name: str | None, num_to_get: int = 25 +) -> Generator[Any]: + db = sqlite3.connect(trace_file) + cur = db.cursor() + limit = num_to_get + if class_name is not None: + cursor = cur.execute( + "SELECT * FROM function_calls WHERE function_name = ? AND file_namename = ? AND class_name = ? ORDER BY time_ns ASC LIMIT ?", + (function_name, file_name, class_name, limit), + ) + else: + cursor = cur.execute( + "SELECT * FROM function_calls WHERE function_name = ? AND file_namename = ? ORDER BY time_ns ASC LIMIT ?", + (function_name, file_name, limit), + ) + + while (val := cursor.fetchone()) is not None: + yield val[8], val[9] + + +def get_function_alias(module: str, function_name: str) -> str: + return "_".join(module.split(".")) + "_" + function_name + + +def create_trace_replay_test( + trace_file: str, functions: list[FunctionModules], test_framework: str = "pytest", max_run_count=100 +) -> str: + assert test_framework in ["pytest", "unittest"] + + imports = f"""import dill as pickle +{"import unittest" if test_framework == "unittest" else ""} +from codeflash.tracing.replay_test import get_next_arg_and_return +""" + + # TODO: Module can have "-" character if the module-root is ".". Need to handle that case + function_properties: list[FunctionProperties] = [ + inspect_top_level_functions_or_methods( + file_name=function.file_name, + function_or_method_name=function.function_name, + class_name=function.class_name, + line_no=function.line_no, + ) + for function in functions + ] + function_imports = [] + for function, function_property in zip(functions, function_properties): + if not function_property.is_top_level: + # can't be imported and run in the replay test + continue + if function_property.is_staticmethod: + function_imports.append( + f"from {function.module_name} import {function_property.staticmethod_class_name} as {get_function_alias(function.module_name, function_property.staticmethod_class_name)}" + ) + elif function.class_name: + function_imports.append( + f"from {function.module_name} import {function.class_name} as {get_function_alias(function.module_name, function.class_name)}" + ) + else: + function_imports.append( + f"from {function.module_name} import {function.function_name} as {get_function_alias(function.module_name, function.function_name)}" + ) + + imports += "\n".join(function_imports) + functions_to_optimize = [function.function_name for function in functions if function.function_name != "__init__"] + metadata = f"""functions = {functions_to_optimize} +trace_file_path = r"{trace_file}" +""" # trace_file_path path is parsed with regex later, format is important + test_function_body = textwrap.dedent( + """\ + for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): + args = pickle.loads(arg_val_pkl) + ret = {function_name}({args}) + """ + ) + test_class_method_body = textwrap.dedent( + """\ + for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + args = pickle.loads(arg_val_pkl){filter_variables} + ret = {class_name_alias}{method_name}(**args) + """ + ) + test_class_staticmethod_body = textwrap.dedent( + """\ + for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): + args = pickle.loads(arg_val_pkl){filter_variables} + ret = {class_name_alias}{method_name}(**args) + """ + ) + if test_framework == "unittest": + self = "self" + test_template = "\nclass TestTracedFunctions(unittest.TestCase):\n" + else: + test_template = "" + self = "" + for func, func_property in zip(functions, function_properties): + if not func_property.is_top_level: + # can't be imported and run in the replay test + continue + if func.class_name is None and not func_property.is_staticmethod: + alias = get_function_alias(func.module_name, func.function_name) + test_body = test_function_body.format( + function_name=alias, + file_name=func.file_name, + orig_function_name=func.function_name, + max_run_count=max_run_count, + args="**args" if func_property.has_args else "", + ) + elif func_property.is_staticmethod: + class_name_alias = get_function_alias(func.module_name, func_property.staticmethod_class_name) + alias = get_function_alias( + func.module_name, func_property.staticmethod_class_name + "_" + func.function_name + ) + method_name = "." + func.function_name if func.function_name != "__init__" else "" + test_body = test_class_staticmethod_body.format( + orig_function_name=func.function_name, + file_name=func.file_name, + class_name_alias=class_name_alias, + method_name=method_name, + max_run_count=max_run_count, + filter_variables="", + ) + else: + class_name_alias = get_function_alias(func.module_name, func.class_name) + alias = get_function_alias(func.module_name, func.class_name + "_" + func.function_name) + + if func_property.is_classmethod: + filter_variables = '\n args.pop("cls", None)' + elif func.function_name == "__init__": + filter_variables = '\n args.pop("__class__", None)' + else: + filter_variables = "" + method_name = "." + func.function_name if func.function_name != "__init__" else "" + test_body = test_class_method_body.format( + orig_function_name=func.function_name, + file_name=func.file_name, + class_name_alias=class_name_alias, + class_name=func.class_name, + method_name=method_name, + max_run_count=max_run_count, + filter_variables=filter_variables, + ) + formatted_test_body = textwrap.indent(test_body, " " if test_framework == "unittest" else " ") + + test_template += " " if test_framework == "unittest" else "" + test_template += f"def test_{alias}({self}):\n{formatted_test_body}\n" + + return imports + "\n" + metadata + "\n" + test_template diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index bec5a03d4..54e0b5118 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -3,12 +3,13 @@ from pathlib import Path import subprocess -def trace_benchmarks_pytest(benchmarks_root: Path, project_root: Path) -> None: +def trace_benchmarks_pytest(benchmarks_root: Path, project_root: Path, output_file: Path) -> None: result = subprocess.run( [ SAFE_SYS_EXECUTABLE, Path(__file__).parent / "pytest_new_process_trace_benchmarks.py", benchmarks_root, + output_file, ], cwd=project_root, check=False, diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 8d28e8c9d..5f5f0ec2f 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -25,7 +25,10 @@ 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 codeflash.benchmarking.codeflash_trace import codeflash_trace + from collections import defaultdict + if TYPE_CHECKING: from argparse import Namespace @@ -104,12 +107,16 @@ def run(self) -> None: #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") + except Exception as e: + logger.info(f"Error while tracing existing benchmarks: {e}") + logger.info(f"Information on existing benchmarks will not be available for this run.") 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]) + codeflash_trace.print_trace_info() # 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) diff --git a/tests/test_codeflash_trace_decorator.py b/tests/test_codeflash_trace_decorator.py index 251b668ec..37234d85a 100644 --- a/tests/test_codeflash_trace_decorator.py +++ b/tests/test_codeflash_trace_decorator.py @@ -1,7 +1,8 @@ from codeflash.benchmarking.codeflash_trace import codeflash_trace from pathlib import Path +from codeflash.code_utils.code_utils import get_run_tmp_file -@codeflash_trace("test_codeflash_trace.trace") +@codeflash_trace def example_function(arr): arr.sort() return arr @@ -12,4 +13,3 @@ def test_codeflash_trace_decorator(): result = example_function(arr) # cleanup test trace file using Path assert result == [1, 2, 3] - Path("test_codeflash_trace.trace").unlink() \ No newline at end of file diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 688e23d8c..071535c6a 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,8 +1,12 @@ from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from pathlib import Path +from codeflash.code_utils.code_utils import get_run_tmp_file 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) \ No newline at end of file + output_file = Path("test_trace_benchmarks.trace").resolve() + trace_benchmarks_pytest(benchmarks_root, project_root, output_file) + assert output_file.exists() + output_file.unlink() \ No newline at end of file From 26b2c4fa6121ba6c9c35ae6335d4d003d6239877 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 13 Mar 2025 18:14:38 -0700 Subject: [PATCH 061/122] initial replay test version working --- code_to_optimize/bubble_sort.py | 2 - .../benchmarks/test_benchmark_bubble_sort.py | 2 +- .../benchmarks/test_process_and_sort.py | 4 +- codeflash/benchmarking/codeflash_trace.py | 122 +++- .../pytest_new_process_trace_benchmarks.py | 9 +- codeflash/benchmarking/replay_test.py | 149 +++-- codeflash/benchmarking/trace_benchmarks.py | 3 +- codeflash/optimization/function_optimizer.py | 15 +- codeflash/tracer.py | 541 +++++++++++------- tests/test_trace_benchmarks.py | 21 +- 10 files changed, 538 insertions(+), 330 deletions(-) diff --git a/code_to_optimize/bubble_sort.py b/code_to_optimize/bubble_sort.py index 91b77f50c..db7db5f92 100644 --- a/code_to_optimize/bubble_sort.py +++ b/code_to_optimize/bubble_sort.py @@ -1,5 +1,3 @@ -from codeflash.benchmarking.codeflash_trace import codeflash_trace -@codeflash_trace def sorter(arr): for i in range(len(arr)): for j in range(len(arr) - 1): diff --git a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py index f1ebcf5c7..21c2bbb29 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py @@ -1,6 +1,6 @@ import pytest -from code_to_optimize.bubble_sort import sorter +from code_to_optimize.bubble_sort_codeflash_trace import sorter def test_sort(benchmark): diff --git a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py index 93d78afef..2713721e4 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py @@ -1,5 +1,5 @@ -from code_to_optimize.process_and_bubble_sort import compute_and_sort -from code_to_optimize.bubble_sort import sorter +from code_to_optimize.process_and_bubble_sort_codeflash_trace import compute_and_sort +from code_to_optimize.bubble_sort_codeflash_trace import sorter def test_compute_and_sort(benchmark): result = benchmark(compute_and_sort, list(reversed(range(5000)))) assert result == 6247083.5 diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 45c7fa6c2..c678b7643 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -3,9 +3,12 @@ import pickle import sqlite3 import time +from pathlib import Path from typing import Callable + + class CodeflashTrace: """A class that provides both a decorator for tracing function calls and a context manager for managing the tracing data lifecycle. @@ -60,8 +63,12 @@ def wrapper(*args, **kwargs): overhead_end_time = time.time() overhead_time = overhead_end_time - overhead_start_time + class_name = "" + qualname = func.__qualname__ + if "." in qualname: + class_name = qualname.split(".")[0] self.function_calls_data.append( - (func.__name__, func.__module__, func.__code__.co_filename, + (func.__name__, class_name, func.__module__, func.__code__.co_filename, benchmark_function_name, benchmark_file_name, execution_time, overhead_time, pickled_args, pickled_kwargs) ) @@ -92,7 +99,7 @@ def write_to_db(self, output_file: str) -> None: # 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, " + "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," "benchmark_function_name TEXT, benchmark_file_name TEXT, " "time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" ) @@ -100,9 +107,9 @@ def write_to_db(self, output_file: str) -> None: # Insert all data at once cur.executemany( "INSERT INTO function_calls " - "(function_name, class_name, file_name, benchmark_function_name, " + "(function_name, class_name, module_name, file_name, benchmark_function_name, " "benchmark_file_name, time_ns, overhead_time_ns, args, kwargs) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", self.function_calls_data ) @@ -153,24 +160,25 @@ def print_codeflash_db(self, limit: int = None) -> None: for i, row in enumerate(cursor.fetchall()): print(f"Record #{i + 1}:") print(f" Function: {row[0]}") - print(f" Module: {row[1]}") - print(f" File: {row[2]}") - print(f" Benchmark Function: {row[3] or 'N/A'}") - print(f" Benchmark File: {row[4] or 'N/A'}") - print(f" Execution Time: {row[5]:.6f} seconds") - print(f" Overhead Time: {row[6]:.6f} seconds") + print(f" Class: {row[1]}") + print(f" Module: {row[2]}") + print(f" File: {row[3]}") + print(f" Benchmark Function: {row[4] or 'N/A'}") + print(f" Benchmark File: {row[5] or 'N/A'}") + print(f" Execution Time: {row[6]:.6f} seconds") + print(f" Overhead Time: {row[7]:.6f} seconds") # Unpickle and print args and kwargs try: - args = pickle.loads(row[7]) - kwargs = pickle.loads(row[8]) + args = pickle.loads(row[8]) + kwargs = pickle.loads(row[9]) print(f" Args: {args}") print(f" Kwargs: {kwargs}") except Exception as e: print(f" Error unpickling args/kwargs: {e}") - print(f" Raw args: {row[7]}") - print(f" Raw kwargs: {row[8]}") + print(f" Raw args: {row[8]}") + print(f" Raw kwargs: {row[9]}") print("\n" + "-" * 40 + "\n") @@ -178,5 +186,91 @@ def print_codeflash_db(self, limit: int = None) -> None: except Exception as e: print(f"Error reading database: {e}") + + def generate_replay_test(self, output_dir: str = None, project_root: str = "", test_framework: str = "pytest", + max_run_count: int = 100) -> None: + """ + Generate multiple replay tests from the traced function calls, grouping by benchmark name. + + Args: + output_dir: Directory to write the generated tests (if None, only returns the code) + project_root: Root directory of the project for module imports + test_framework: 'pytest' or 'unittest' + max_run_count: Maximum number of runs to include per function + + Returns: + Dictionary mapping benchmark names to generated test code + """ + import isort + from codeflash.verification.verification_utils import get_test_file_path + + if not self.db_path: + print("No database path set. Call write_to_db first or set db_path manually.") + return {} + + try: + # Import the function here to avoid circular imports + from codeflash.benchmarking.replay_test import create_trace_replay_test + + print("connecting to: ", self.db_path) + # Connect to the database + conn = sqlite3.connect(self.db_path) + cursor = conn.cursor() + + # Get distinct benchmark names + cursor.execute( + "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM function_calls" + ) + benchmarks = cursor.fetchall() + + # Generate a test for each benchmark + for benchmark in benchmarks: + benchmark_function_name, benchmark_file_name = benchmark + # Get functions associated with this benchmark + cursor.execute( + "SELECT DISTINCT function_name, class_name, module_name, file_name FROM function_calls " + "WHERE benchmark_function_name = ? AND benchmark_file_name = ?", + (benchmark_function_name, benchmark_file_name) + ) + + functions_data = [] + for func_row in cursor.fetchall(): + function_name, class_name, module_name, file_name = func_row + + # Add this function to our list + functions_data.append({ + "function_name": function_name, + "class_name": class_name, + "file_name": file_name, + "module_name": module_name + }) + + if not functions_data: + print(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") + continue + + # Generate the test code for this benchmark + test_code = create_trace_replay_test( + trace_file=self.db_path, + functions_data=functions_data, + test_framework=test_framework, + max_run_count=max_run_count, + ) + test_code = isort.code(test_code) + + # Write to file if requested + if output_dir: + output_file = get_test_file_path( + test_dir=Path(output_dir), function_name=f"{benchmark_file_name[5:]}_{benchmark_function_name}", test_type="replay" + ) + with open(output_file, 'w') as f: + f.write(test_code) + print(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") + + conn.close() + + except Exception as e: + print(f"Error generating replay tests: {e}") + # Create a singleton instance codeflash_trace = CodeflashTrace() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index ebe1fa4ae..8e1958fec 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -1,10 +1,16 @@ import sys +from pathlib import Path + +from codeflash.verification.verification_utils import get_test_file_path from plugin.plugin import CodeFlashPlugin from codeflash.benchmarking.codeflash_trace import codeflash_trace from codeflash.code_utils.code_utils import get_run_tmp_file benchmarks_root = sys.argv[1] -output_file = sys.argv[2] +tests_root = sys.argv[2] +output_file = sys.argv[3] +# current working directory +project_root = Path.cwd() if __name__ == "__main__": import pytest @@ -14,6 +20,7 @@ ) codeflash_trace.write_to_db(output_file) codeflash_trace.print_codeflash_db() + codeflash_trace.generate_replay_test(tests_root, project_root, test_framework="pytest") except Exception as e: print(f"Failed to collect tests: {e!s}") exitcode = -1 \ No newline at end of file diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 0bc2de1d4..9bc2c79f3 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -3,31 +3,29 @@ import sqlite3 import textwrap from collections.abc import Generator -from typing import Any, List, Optional - -from codeflash.discovery.functions_to_optimize import FunctionProperties, inspect_top_level_functions_or_methods -from codeflash.tracing.tracing_utils import FunctionModules +from typing import Any, Dict def get_next_arg_and_return( - trace_file: str, function_name: str, file_name: str, class_name: str | None, num_to_get: int = 25 + trace_file: str, function_name: str, file_name: str, class_name: str | None = None, num_to_get: int = 25 ) -> Generator[Any]: db = sqlite3.connect(trace_file) cur = db.cursor() limit = num_to_get + if class_name is not None: cursor = cur.execute( - "SELECT * FROM function_calls WHERE function_name = ? AND file_namename = ? AND class_name = ? ORDER BY time_ns ASC LIMIT ?", + "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? AND class_name = ? ORDER BY time_ns ASC LIMIT ?", (function_name, file_name, class_name, limit), ) else: cursor = cur.execute( - "SELECT * FROM function_calls WHERE function_name = ? AND file_namename = ? ORDER BY time_ns ASC LIMIT ?", + "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? ORDER BY time_ns ASC LIMIT ?", (function_name, file_name, limit), ) while (val := cursor.fetchone()) is not None: - yield val[8], val[9] + yield val[8], val[9] # args and kwargs are at indices 7 and 8 def get_function_alias(module: str, function_name: str) -> str: @@ -35,122 +33,109 @@ def get_function_alias(module: str, function_name: str) -> str: def create_trace_replay_test( - trace_file: str, functions: list[FunctionModules], test_framework: str = "pytest", max_run_count=100 + trace_file: str, + functions_data: list[Dict[str, Any]], + test_framework: str = "pytest", + max_run_count=100 ) -> str: + """Create a replay test for functions based on trace data. + + Args: + trace_file: Path to the SQLite database file + functions_data: List of dictionaries with function info extracted from DB + test_framework: 'pytest' or 'unittest' + max_run_count: Maximum number of runs to include in the test + + Returns: + A string containing the test code + + """ assert test_framework in ["pytest", "unittest"] imports = f"""import dill as pickle {"import unittest" if test_framework == "unittest" else ""} -from codeflash.tracing.replay_test import get_next_arg_and_return +from codeflash.benchmarking.replay_test import get_next_arg_and_return """ - # TODO: Module can have "-" character if the module-root is ".". Need to handle that case - function_properties: list[FunctionProperties] = [ - inspect_top_level_functions_or_methods( - file_name=function.file_name, - function_or_method_name=function.function_name, - class_name=function.class_name, - line_no=function.line_no, - ) - for function in functions - ] function_imports = [] - for function, function_property in zip(functions, function_properties): - if not function_property.is_top_level: - # can't be imported and run in the replay test - continue - if function_property.is_staticmethod: - function_imports.append( - f"from {function.module_name} import {function_property.staticmethod_class_name} as {get_function_alias(function.module_name, function_property.staticmethod_class_name)}" - ) - elif function.class_name: + for func in functions_data: + module_name = func.get("module_name") + function_name = func.get("function_name") + class_name = func.get("class_name", "") + + if class_name: function_imports.append( - f"from {function.module_name} import {function.class_name} as {get_function_alias(function.module_name, function.class_name)}" + f"from {module_name} import {class_name} as {get_function_alias(module_name, class_name)}" ) else: function_imports.append( - f"from {function.module_name} import {function.function_name} as {get_function_alias(function.module_name, function.function_name)}" + f"from {module_name} import {function_name} as {get_function_alias(module_name, function_name)}" ) imports += "\n".join(function_imports) - functions_to_optimize = [function.function_name for function in functions if function.function_name != "__init__"] + + functions_to_optimize = [func.get("function_name") for func in functions_data + if func.get("function_name") != "__init__"] metadata = f"""functions = {functions_to_optimize} trace_file_path = r"{trace_file}" -""" # trace_file_path path is parsed with regex later, format is important +""" + + # Templates for different types of tests test_function_body = textwrap.dedent( """\ - for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): - args = pickle.loads(arg_val_pkl) - ret = {function_name}({args}) + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + ret = {function_name}(*args, **kwargs) """ ) + test_class_method_body = textwrap.dedent( """\ - for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): - args = pickle.loads(arg_val_pkl){filter_variables} - ret = {class_name_alias}{method_name}(**args) - """ - ) - test_class_staticmethod_body = textwrap.dedent( - """\ - for arg_val_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): - args = pickle.loads(arg_val_pkl){filter_variables} - ret = {class_name_alias}{method_name}(**args) + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl){filter_variables} + ret = {class_name_alias}{method_name}(**args, **kwargs) """ ) + if test_framework == "unittest": self = "self" test_template = "\nclass TestTracedFunctions(unittest.TestCase):\n" else: test_template = "" self = "" - for func, func_property in zip(functions, function_properties): - if not func_property.is_top_level: - # can't be imported and run in the replay test - continue - if func.class_name is None and not func_property.is_staticmethod: - alias = get_function_alias(func.module_name, func.function_name) + + for func in functions_data: + module_name = func.get("module_name") + function_name = func.get("function_name") + class_name = func.get("class_name") + file_name = func.get("file_name") + + if not class_name: + alias = get_function_alias(module_name, function_name) test_body = test_function_body.format( function_name=alias, - file_name=func.file_name, - orig_function_name=func.function_name, + file_name=file_name, + orig_function_name=function_name, max_run_count=max_run_count, - args="**args" if func_property.has_args else "", - ) - elif func_property.is_staticmethod: - class_name_alias = get_function_alias(func.module_name, func_property.staticmethod_class_name) - alias = get_function_alias( - func.module_name, func_property.staticmethod_class_name + "_" + func.function_name - ) - method_name = "." + func.function_name if func.function_name != "__init__" else "" - test_body = test_class_staticmethod_body.format( - orig_function_name=func.function_name, - file_name=func.file_name, - class_name_alias=class_name_alias, - method_name=method_name, - max_run_count=max_run_count, - filter_variables="", ) else: - class_name_alias = get_function_alias(func.module_name, func.class_name) - alias = get_function_alias(func.module_name, func.class_name + "_" + func.function_name) - - if func_property.is_classmethod: - filter_variables = '\n args.pop("cls", None)' - elif func.function_name == "__init__": - filter_variables = '\n args.pop("__class__", None)' - else: - filter_variables = "" - method_name = "." + func.function_name if func.function_name != "__init__" else "" + class_name_alias = get_function_alias(module_name, class_name) + alias = get_function_alias(module_name, class_name + "_" + function_name) + + filter_variables = "" + method_name = "." + function_name if function_name != "__init__" else "" test_body = test_class_method_body.format( - orig_function_name=func.function_name, - file_name=func.file_name, + orig_function_name=function_name, + file_name=file_name, class_name_alias=class_name_alias, - class_name=func.class_name, + class_name=class_name, method_name=method_name, max_run_count=max_run_count, filter_variables=filter_variables, ) + formatted_test_body = textwrap.indent(test_body, " " if test_framework == "unittest" else " ") test_template += " " if test_framework == "unittest" else "" diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 54e0b5118..5c0a077dc 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -3,12 +3,13 @@ from pathlib import Path import subprocess -def trace_benchmarks_pytest(benchmarks_root: Path, project_root: Path, output_file: Path) -> None: +def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, output_file: Path) -> None: result = subprocess.run( [ SAFE_SYS_EXECUTABLE, Path(__file__).parent / "pytest_new_process_trace_benchmarks.py", benchmarks_root, + tests_root, output_file, ], cwd=project_root, diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 38277851b..249251e34 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -140,19 +140,12 @@ def optimize_function(self) -> Result[BestOptimization, str]: with helper_function_path.open(encoding="utf8") as f: helper_code = f.read() original_helper_code[helper_function_path] = helper_code + if has_any_async_functions(code_context.read_writable_code): + return Failure("Codeflash does not support async functions in the code to optimize.") - logger.info("Code to be optimized:") code_print(code_context.read_writable_code) - - for module_abspath, helper_code_source in original_helper_code.items(): - code_context.code_to_optimize_with_helpers = add_needed_imports_from_module( - helper_code_source, - code_context.code_to_optimize_with_helpers, - module_abspath, - self.function_to_optimize.file_path, - self.args.project_root, - ) - + logger.info("Read only code") + code_print(code_context.read_only_context_code) generated_test_paths = [ get_test_file_path( self.test_cfg.tests_root, self.function_to_optimize.function_name, test_index, test_type="unit" diff --git a/codeflash/tracer.py b/codeflash/tracer.py index 02a0e4157..eb4df84d4 100644 --- a/codeflash/tracer.py +++ b/codeflash/tracer.py @@ -18,19 +18,21 @@ import os import pathlib import pickle -import re import sqlite3 import sys +import threading import time +from argparse import ArgumentParser from collections import defaultdict -from copy import copy -from io import StringIO from pathlib import Path -from types import FrameType -from typing import Any, ClassVar, List +from typing import TYPE_CHECKING, Any, Callable, ClassVar import dill import isort +from rich.align import Align +from rich.panel import Panel +from rich.table import Table +from rich.text import Text from codeflash.cli_cmds.cli import project_root_from_module_root from codeflash.cli_cmds.console import console @@ -40,14 +42,34 @@ from codeflash.tracing.replay_test import create_trace_replay_test from codeflash.tracing.tracing_utils import FunctionModules from codeflash.verification.verification_utils import get_test_file_path -# import warnings -# warnings.filterwarnings("ignore", category=dill.PickleWarning) -# warnings.filterwarnings("ignore", category=DeprecationWarning) + +if TYPE_CHECKING: + from types import FrameType, TracebackType + + +class FakeCode: + def __init__(self, filename: str, line: int, name: str) -> None: + self.co_filename = filename + self.co_line = line + self.co_name = name + self.co_firstlineno = 0 + + def __repr__(self) -> str: + return repr((self.co_filename, self.co_line, self.co_name, None)) + + +class FakeFrame: + def __init__(self, code: FakeCode, prior: FakeFrame | None) -> None: + self.f_code = code + self.f_back = prior + self.f_locals: dict = {} + # Debug this file by simply adding print statements. This file is not meant to be debugged by the debugger. class Tracer: - """Use this class as a 'with' context manager to trace a function call, - input arguments, and profiling info. + """Use this class as a 'with' context manager to trace a function call. + + Traces function calls, input arguments, and profiling info. """ def __init__( @@ -58,9 +80,10 @@ def __init__( config_file_path: Path | None = None, max_function_count: int = 256, timeout: int | None = None, # seconds - benchmark: bool = False, ) -> None: - """:param output: The path to the output trace file + """Use this class to trace function calls. + + :param output: The path to the output trace file :param functions: List of functions to trace. If None, trace all functions :param disable: Disable the tracer if True :param config_file_path: Path to the pyproject.toml file, if None then it will be auto-discovered @@ -71,7 +94,9 @@ def __init__( if functions is None: functions = [] if os.environ.get("CODEFLASH_TRACER_DISABLE", "0") == "1": - console.print("Codeflash: Tracer disabled by environment variable CODEFLASH_TRACER_DISABLE") + console.rule( + "Codeflash: Tracer disabled by environment variable CODEFLASH_TRACER_DISABLE", style="bold red" + ) disable = True self.disable = disable if self.disable: @@ -86,7 +111,7 @@ def __init__( self.con = None self.output_file = Path(output).resolve() self.functions = functions - self.function_modules: List[FunctionModules] = [] + self.function_modules: list[FunctionModules] = [] self.function_count = defaultdict(int) self.current_file_path = Path(__file__).resolve() self.ignored_qualified_functions = { @@ -96,16 +121,16 @@ def __init__( self.max_function_count = max_function_count self.config, found_config_path = parse_config_file(config_file_path) self.project_root = project_root_from_module_root(Path(self.config["module_root"]), found_config_path) + console.rule(f"Project Root: {self.project_root}", style="bold blue") self.ignored_functions = {"", "", "", "", "", ""} - self.file_being_called_from: str = str(Path(sys._getframe().f_back.f_code.co_filename).name).replace(".", "_") + self.file_being_called_from: str = str(Path(sys._getframe().f_back.f_code.co_filename).name).replace(".", "_") # noqa: SLF001 assert timeout is None or timeout > 0, "Timeout should be greater than 0" self.timeout = timeout self.next_insert = 1000 self.trace_count = 0 - self.benchmark = benchmark # Profiler variables self.bias = 0 # calibration constant self.timings = {} @@ -120,48 +145,44 @@ def __init__( def __enter__(self) -> None: if self.disable: return - - # if getattr(Tracer, "used_once", False): - # console.print( - # "Codeflash: Tracer can only be used once per program run. " - # "Please only enable the Tracer once. Skipping tracing this section." - # ) - # self.disable = True - # return - # Tracer.used_once = True + if getattr(Tracer, "used_once", False): + console.print( + "Codeflash: Tracer can only be used once per program run. " + "Please only enable the Tracer once. Skipping tracing this section." + ) + self.disable = True + return + Tracer.used_once = True if pathlib.Path(self.output_file).exists(): - console.print("Codeflash: Removing existing trace file") + console.rule("Removing existing trace file", style="bold red") + console.rule() pathlib.Path(self.output_file).unlink(missing_ok=True) - self.con = sqlite3.connect(self.output_file) + self.con = sqlite3.connect(self.output_file, check_same_thread=False) cur = self.con.cursor() cur.execute("""PRAGMA synchronous = OFF""") + cur.execute("""PRAGMA journal_mode = WAL""") # TODO: Check out if we need to export the function test name as well cur.execute( "CREATE TABLE function_calls(type TEXT, function TEXT, classname TEXT, filename TEXT, " "line_number INTEGER, last_frame_address INTEGER, time_ns INTEGER, args BLOB)" ) - console.print("Codeflash: Tracing started!") - frame = sys._getframe(0) # Get this frame and simulate a call to it + console.rule("Codeflash: Traced Program Output Begin", style="bold blue") + frame = sys._getframe(0) # Get this frame and simulate a call to it # noqa: SLF001 self.dispatch["call"](self, frame, 0) self.start_time = time.time() sys.setprofile(self.trace_callback) + threading.setprofile(self.trace_callback) - def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + def __exit__( + self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None + ) -> None: if self.disable: return sys.setprofile(None) self.con.commit() - # Check if any functions were actually traced - if self.trace_count == 0: - self.con.close() - # Delete the trace file if no functions were traced - if self.output_file.exists(): - self.output_file.unlink() - console.print("Codeflash: No functions were traced. Removing trace database.") - return - + console.rule("Codeflash: Traced Program Output End", style="bold blue") self.create_stats() cur = self.con.cursor() @@ -185,25 +206,18 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: cur.execute("INSERT INTO total_time VALUES (?)", (self.total_tt,)) self.con.commit() self.con.close() - function_string = [str(function.file_name) + ":" + (function.class_name + ":" if function.class_name else "") + function.function_name for function in self.function_modules] - # print(function_string) # filter any functions where we did not capture the return - # self.function_modules = [ - # function - # for function in self.function_modules - # if self.function_count[ - # str(function.file_name) - # + ":" - # + (function.class_name + ":" if function.class_name else "") - # + function.function_name - # ] - # > 0 - # ] self.function_modules = [ function for function in self.function_modules - if str(str(function.file_name) + ":" + (function.class_name + ":" if function.class_name else "") + function.function_name) in self.function_count + if self.function_count[ + str(function.file_name) + + ":" + + (function.class_name + ":" if function.class_name else "") + + function.function_name + ] + > 0 ] replay_test = create_trace_replay_test( @@ -212,24 +226,15 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: test_framework=self.config["test_framework"], max_run_count=self.max_function_count, ) - # Need a better way to store the replay test - # function_path = "_".join(self.functions) if self.functions else self.file_being_called_from - function_path = self.file_being_called_from - if self.benchmark and self.config["benchmarks_root"]: - # check if replay test dir exists, create - replay_test_dir = Path(self.config["benchmarks_root"]) / "codeflash_replay_tests" - if not replay_test_dir.exists(): - replay_test_dir.mkdir(parents=True) - test_file_path = get_test_file_path( - test_dir=replay_test_dir, function_name=function_path, test_type="replay" - ) - else: - test_file_path = get_test_file_path( - test_dir=Path(self.config["tests_root"]), function_name=function_path, test_type="replay" - ) + function_path = "_".join(self.functions) if self.functions else self.file_being_called_from + test_file_path = get_test_file_path( + test_dir=Path(self.config["tests_root"]), function_name=function_path, test_type="replay" + ) replay_test = isort.code(replay_test) - with open(test_file_path, "w", encoding="utf8") as file: + + with Path(test_file_path).open("w", encoding="utf8") as file: file.write(replay_test) + console.print( f"Codeflash: Traced {self.trace_count} function calls successfully and replay test created at - {test_file_path}", crop=False, @@ -237,14 +242,13 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: overflow="ignore", ) - def tracer_logic(self, frame: FrameType, event: str): + def tracer_logic(self, frame: FrameType, event: str) -> None: if event != "call": return - if self.timeout is not None: - if (time.time() - self.start_time) > self.timeout: - sys.setprofile(None) - console.print(f"Codeflash: Timeout reached! Stopping tracing at {self.timeout} seconds.") - return + if self.timeout is not None and (time.time() - self.start_time) > self.timeout: + sys.setprofile(None) + console.print(f"Codeflash: Timeout reached! Stopping tracing at {self.timeout} seconds.") + return code = frame.f_code file_name = Path(code.co_filename).resolve() # TODO : It currently doesn't log the last return call from the first function @@ -253,9 +257,8 @@ def tracer_logic(self, frame: FrameType, event: str): return if not file_name.exists(): return - # if self.functions: - # if code.co_name not in self.functions: - # return + if self.functions and code.co_name not in self.functions: + return class_name = None arguments = frame.f_locals try: @@ -267,16 +270,12 @@ def tracer_logic(self, frame: FrameType, event: str): class_name = arguments["self"].__class__.__name__ elif "cls" in arguments and hasattr(arguments["cls"], "__name__"): class_name = arguments["cls"].__name__ - except: + except: # noqa: E722 # someone can override the getattr method and raise an exception. I'm looking at you wrapt return - function_qualified_name = f"{file_name}:{(class_name + ':' if class_name else '')}{code.co_name}" if function_qualified_name in self.ignored_qualified_functions: return - if self.functions and function_qualified_name not in self.functions: - return - if function_qualified_name not in self.function_count: # seeing this function for the first time self.function_count[function_qualified_name] = 0 @@ -351,17 +350,14 @@ def tracer_logic(self, frame: FrameType, event: str): self.next_insert = 1000 self.con.commit() - def trace_callback(self, frame: FrameType, event: str, arg: Any) -> None: + def trace_callback(self, frame: FrameType, event: str, arg: str | None) -> None: # profiler section timer = self.timer t = timer() - self.t - self.bias if event == "c_call": self.c_func_name = arg.__name__ - if self.dispatch[event](self, frame, t): - prof_success = True - else: - prof_success = False + prof_success = bool(self.dispatch[event](self, frame, t)) # tracer section self.tracer_logic(frame, event) # measure the time as the last thing before return @@ -370,45 +366,60 @@ def trace_callback(self, frame: FrameType, event: str, arg: Any) -> None: else: self.t = timer() - t # put back unrecorded delta - def trace_dispatch_call(self, frame, t): - if self.cur and frame.f_back is not self.cur[-2]: - rpt, rit, ret, rfn, rframe, rcur = self.cur - if not isinstance(rframe, Tracer.fake_frame): - assert rframe.f_back is frame.f_back, ("Bad call", rfn, rframe, rframe.f_back, frame, frame.f_back) - self.trace_dispatch_return(rframe, 0) - assert self.cur is None or frame.f_back is self.cur[-2], ("Bad call", self.cur[-3]) - fcode = frame.f_code - arguments = frame.f_locals - class_name = None + def trace_dispatch_call(self, frame: FrameType, t: int) -> int: + """Handle call events in the profiler.""" try: - if ( - "self" in arguments - and hasattr(arguments["self"], "__class__") - and hasattr(arguments["self"].__class__, "__name__") - ): - class_name = arguments["self"].__class__.__name__ - elif "cls" in arguments and hasattr(arguments["cls"], "__name__"): - class_name = arguments["cls"].__name__ - except: - pass - fn = (fcode.co_filename, fcode.co_firstlineno, fcode.co_name, class_name) - self.cur = (t, 0, 0, fn, frame, self.cur) - timings = self.timings - if fn in timings: - cc, ns, tt, ct, callers = timings[fn] - timings[fn] = cc, ns + 1, tt, ct, callers - else: - timings[fn] = 0, 0, 0, 0, {} - return 1 + # In multi-threaded contexts, we need to be more careful about frame comparisons + if self.cur and frame.f_back is not self.cur[-2]: + # This happens when we're in a different thread + rpt, rit, ret, rfn, rframe, rcur = self.cur + + # Only attempt to handle the frame mismatch if we have a valid rframe + if ( + not isinstance(rframe, FakeFrame) + and hasattr(rframe, "f_back") + and hasattr(frame, "f_back") + and rframe.f_back is frame.f_back + ): + self.trace_dispatch_return(rframe, 0) + + # Get function information + fcode = frame.f_code + arguments = frame.f_locals + class_name = None + try: + if ( + "self" in arguments + and hasattr(arguments["self"], "__class__") + and hasattr(arguments["self"].__class__, "__name__") + ): + class_name = arguments["self"].__class__.__name__ + elif "cls" in arguments and hasattr(arguments["cls"], "__name__"): + class_name = arguments["cls"].__name__ + except Exception: # noqa: BLE001, S110 + pass + + fn = (fcode.co_filename, fcode.co_firstlineno, fcode.co_name, class_name) + self.cur = (t, 0, 0, fn, frame, self.cur) + timings = self.timings + if fn in timings: + cc, ns, tt, ct, callers = timings[fn] + timings[fn] = cc, ns + 1, tt, ct, callers + else: + timings[fn] = 0, 0, 0, 0, {} + return 1 # noqa: TRY300 + except Exception: # noqa: BLE001 + # Handle any errors gracefully + return 0 - def trace_dispatch_exception(self, frame, t): + def trace_dispatch_exception(self, frame: FrameType, t: int) -> int: rpt, rit, ret, rfn, rframe, rcur = self.cur if (rframe is not frame) and rcur: return self.trace_dispatch_return(rframe, t) self.cur = rpt, rit + t, ret, rfn, rframe, rcur return 1 - def trace_dispatch_c_call(self, frame, t): + def trace_dispatch_c_call(self, frame: FrameType, t: int) -> int: fn = ("", 0, self.c_func_name, None) self.cur = (t, 0, 0, fn, frame, self.cur) timings = self.timings @@ -419,15 +430,27 @@ def trace_dispatch_c_call(self, frame, t): timings[fn] = 0, 0, 0, 0, {} return 1 - def trace_dispatch_return(self, frame, t): - if frame is not self.cur[-2]: - assert frame is self.cur[-2].f_back, ("Bad return", self.cur[-3]) - self.trace_dispatch_return(self.cur[-2], 0) + def trace_dispatch_return(self, frame: FrameType, t: int) -> int: + if not self.cur or not self.cur[-2]: + return 0 + # In multi-threaded environments, frames can get mismatched + if frame is not self.cur[-2]: + # Don't assert in threaded environments - frames can legitimately differ + if hasattr(frame, "f_back") and hasattr(self.cur[-2], "f_back") and frame is self.cur[-2].f_back: + self.trace_dispatch_return(self.cur[-2], 0) + else: + # We're in a different thread or context, can't continue with this frame + return 0 # Prefix "r" means part of the Returning or exiting frame. # Prefix "p" means part of the Previous or Parent or older frame. rpt, rit, ret, rfn, frame, rcur = self.cur + + # Guard against invalid rcur (w threading) + if not rcur: + return 0 + rit = rit + t frame_total = rit + ret @@ -435,6 +458,9 @@ def trace_dispatch_return(self, frame, t): self.cur = ppt, pit + rpt, pet + frame_total, pfn, pframe, pcur timings = self.timings + if rfn not in timings: + # w threading, rfn can be missing + timings[rfn] = 0, 0, 0, 0, {} cc, ns, tt, ct, callers = timings[rfn] if not ns: # This is the only occurrence of the function on the stack. @@ -456,7 +482,7 @@ def trace_dispatch_return(self, frame, t): return 1 - dispatch: ClassVar[dict[str, callable]] = { + dispatch: ClassVar[dict[str, Callable[[Tracer, FrameType, int], int]]] = { "call": trace_dispatch_call, "exception": trace_dispatch_exception, "return": trace_dispatch_return, @@ -465,32 +491,13 @@ def trace_dispatch_return(self, frame, t): "c_return": trace_dispatch_return, } - class fake_code: - def __init__(self, filename, line, name): - self.co_filename = filename - self.co_line = line - self.co_name = name - self.co_firstlineno = 0 - - def __repr__(self): - return repr((self.co_filename, self.co_line, self.co_name, None)) - - class fake_frame: - def __init__(self, code, prior): - self.f_code = code - self.f_back = prior - self.f_locals = {} - - def simulate_call(self, name): - code = self.fake_code("profiler", 0, name) - if self.cur: - pframe = self.cur[-2] - else: - pframe = None - frame = self.fake_frame(code, pframe) + def simulate_call(self, name: str) -> None: + code = FakeCode("profiler", 0, name) + pframe = self.cur[-2] if self.cur else None + frame = FakeFrame(code, pframe) self.dispatch["call"](self, frame, 0) - def simulate_cmd_complete(self): + def simulate_cmd_complete(self) -> None: get_time = self.timer t = get_time() - self.t while self.cur[-1]: @@ -500,60 +507,174 @@ def simulate_cmd_complete(self): t = 0 self.t = get_time() - t - def print_stats(self, sort=-1): - import pstats + def print_stats(self, sort: str | int | tuple = -1) -> None: + if not self.stats: + console.print("Codeflash: No stats available to print") + self.total_tt = 0 + return if not isinstance(sort, tuple): sort = (sort,) - # The following code customizes the default printing behavior to - # print in milliseconds. - s = StringIO() - stats_obj = pstats.Stats(copy(self), stream=s) - stats_obj.strip_dirs().sort_stats(*sort).print_stats(100) - self.total_tt = stats_obj.total_tt - console.print("total_tt", self.total_tt) - raw_stats = s.getvalue() - m = re.search(r"function calls?.*in (\d+)\.\d+ (seconds?)", raw_stats) - total_time = None - if m: - total_time = int(m.group(1)) - if total_time is None: - console.print("Failed to get total time from stats") - total_time_ms = total_time / 1e6 - raw_stats = re.sub( - r"(function calls?.*)in (\d+)\.\d+ (seconds?)", rf"\1 in {total_time_ms:.3f} milliseconds", raw_stats - ) - match_pattern = r"^ *[\d\/]+ +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +" - m = re.findall(match_pattern, raw_stats, re.MULTILINE) - ms_times = [] - for tottime, percall, cumtime, percall_cum in m: - tottime_ms = int(tottime) / 1e6 - percall_ms = int(percall) / 1e6 - cumtime_ms = int(cumtime) / 1e6 - percall_cum_ms = int(percall_cum) / 1e6 - ms_times.append([tottime_ms, percall_ms, cumtime_ms, percall_cum_ms]) - split_stats = raw_stats.split("\n") - new_stats = [] - - replace_pattern = r"^( *[\d\/]+) +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +(\d+)\.\d+ +(.*)" - times_index = 0 - for line in split_stats: - if times_index >= len(ms_times): - replaced = line - else: - replaced, n = re.subn( - replace_pattern, - rf"\g<1>{ms_times[times_index][0]:8.3f} {ms_times[times_index][1]:8.3f} {ms_times[times_index][2]:8.3f} {ms_times[times_index][3]:8.3f} \g<6>", - line, - count=1, + + # First, convert stats to make them pstats-compatible + try: + # Initialize empty collections for pstats + self.files = [] + self.top_level = [] + + # Create entirely new dictionaries instead of modifying existing ones + new_stats = {} + new_timings = {} + + # Convert stats dictionary + stats_items = list(self.stats.items()) + for func, stats_data in stats_items: + try: + # Make sure we have 5 elements in stats_data + if len(stats_data) != 5: + console.print(f"Skipping malformed stats data for {func}: {stats_data}") + continue + + cc, nc, tt, ct, callers = stats_data + + if len(func) == 4: + file_name, line_num, func_name, class_name = func + new_func_name = f"{class_name}.{func_name}" if class_name else func_name + new_func = (file_name, line_num, new_func_name) + else: + new_func = func # Keep as is if already in correct format + + new_callers = {} + callers_items = list(callers.items()) + for caller_func, count in callers_items: + if isinstance(caller_func, tuple): + if len(caller_func) == 4: + caller_file, caller_line, caller_name, caller_class = caller_func + caller_new_name = f"{caller_class}.{caller_name}" if caller_class else caller_name + new_caller_func = (caller_file, caller_line, caller_new_name) + else: + new_caller_func = caller_func + else: + console.print(f"Unexpected caller format: {caller_func}") + new_caller_func = str(caller_func) + + new_callers[new_caller_func] = count + + # Store with new format + new_stats[new_func] = (cc, nc, tt, ct, new_callers) + except Exception as e: # noqa: BLE001 + console.print(f"Error converting stats for {func}: {e}") + continue + + timings_items = list(self.timings.items()) + for func, timing_data in timings_items: + try: + if len(timing_data) != 5: + console.print(f"Skipping malformed timing data for {func}: {timing_data}") + continue + + cc, ns, tt, ct, callers = timing_data + + if len(func) == 4: + file_name, line_num, func_name, class_name = func + new_func_name = f"{class_name}.{func_name}" if class_name else func_name + new_func = (file_name, line_num, new_func_name) + else: + new_func = func + + new_callers = {} + callers_items = list(callers.items()) + for caller_func, count in callers_items: + if isinstance(caller_func, tuple): + if len(caller_func) == 4: + caller_file, caller_line, caller_name, caller_class = caller_func + caller_new_name = f"{caller_class}.{caller_name}" if caller_class else caller_name + new_caller_func = (caller_file, caller_line, caller_new_name) + else: + new_caller_func = caller_func + else: + console.print(f"Unexpected caller format: {caller_func}") + new_caller_func = str(caller_func) + + new_callers[new_caller_func] = count + + new_timings[new_func] = (cc, ns, tt, ct, new_callers) + except Exception as e: # noqa: BLE001 + console.print(f"Error converting timings for {func}: {e}") + continue + + self.stats = new_stats + self.timings = new_timings + + self.total_tt = sum(tt for _, _, tt, _, _ in self.stats.values()) + + total_calls = sum(cc for cc, _, _, _, _ in self.stats.values()) + total_primitive = sum(nc for _, nc, _, _, _ in self.stats.values()) + + summary = Text.assemble( + f"{total_calls:,} function calls ", + ("(" + f"{total_primitive:,} primitive calls" + ")", "dim"), + f" in {self.total_tt / 1e6:.3f}milliseconds", + ) + + console.print(Align.center(Panel(summary, border_style="blue", width=80, padding=(0, 2), expand=False))) + + table = Table( + show_header=True, + header_style="bold magenta", + border_style="blue", + title="[bold]Function Profile[/bold] (ordered by internal time)", + title_style="cyan", + caption=f"Showing top 25 of {len(self.stats)} functions", + ) + + table.add_column("Calls", justify="right", style="green", width=10) + table.add_column("Time (ms)", justify="right", style="cyan", width=10) + table.add_column("Per Call", justify="right", style="cyan", width=10) + table.add_column("Cum (ms)", justify="right", style="yellow", width=10) + table.add_column("Cum/Call", justify="right", style="yellow", width=10) + table.add_column("Function", style="blue") + + sorted_stats = sorted( + ((func, stats) for func, stats in self.stats.items() if isinstance(func, tuple) and len(func) == 3), + key=lambda x: x[1][2], # Sort by tt (internal time) + reverse=True, + )[:25] # Limit to top 25 + + # Format and add each row to the table + for func, (cc, nc, tt, ct, _) in sorted_stats: + filename, lineno, funcname = func + + # Format calls - show recursive format if different + calls_str = f"{cc}/{nc}" if cc != nc else f"{cc:,}" + + # Convert to milliseconds + tt_ms = tt / 1e6 + ct_ms = ct / 1e6 + + # Calculate per-call times + per_call = tt_ms / cc if cc > 0 else 0 + cum_per_call = ct_ms / nc if nc > 0 else 0 + base_filename = Path(filename).name + file_link = f"[link=file://{filename}]{base_filename}[/link]" + + table.add_row( + calls_str, + f"{tt_ms:.3f}", + f"{per_call:.3f}", + f"{ct_ms:.3f}", + f"{cum_per_call:.3f}", + f"{funcname} [dim]({file_link}:{lineno})[/dim]", ) - if n > 0: - times_index += 1 - new_stats.append(replaced) - console.print("\n".join(new_stats)) + console.print(Align.center(table)) + + except Exception as e: # noqa: BLE001 + console.print(f"[bold red]Error in stats processing:[/bold red] {e}") + console.print(f"Traced {self.trace_count:,} function calls") + self.total_tt = 0 - def make_pstats_compatible(self): + def make_pstats_compatible(self) -> None: # delete the extra class_name item from the function tuple self.files = [] self.top_level = [] @@ -568,36 +689,33 @@ def make_pstats_compatible(self): self.stats = new_stats self.timings = new_timings - def dump_stats(self, file): - with open(file, "wb") as f: - self.create_stats() + def dump_stats(self, file: str) -> None: + with Path(file).open("wb") as f: marshal.dump(self.stats, f) - def create_stats(self): + def create_stats(self) -> None: self.simulate_cmd_complete() self.snapshot_stats() - def snapshot_stats(self): + def snapshot_stats(self) -> None: self.stats = {} - for func, (cc, ns, tt, ct, callers) in self.timings.items(): - callers = callers.copy() + for func, (cc, _ns, tt, ct, caller_dict) in self.timings.items(): + callers = caller_dict.copy() nc = 0 for callcnt in callers.values(): nc += callcnt self.stats[func] = cc, nc, tt, ct, callers - def runctx(self, cmd, globals, locals): + def runctx(self, cmd: str, global_vars: dict[str, Any], local_vars: dict[str, Any]) -> Tracer | None: self.__enter__() try: - exec(cmd, globals, locals) + exec(cmd, global_vars, local_vars) # noqa: S102 finally: self.__exit__(None, None, None) return self -def main(): - from argparse import ArgumentParser - +def main() -> ArgumentParser: parser = ArgumentParser(allow_abbrev=False) parser.add_argument("-o", "--outfile", dest="outfile", help="Save trace to ", required=True) parser.add_argument("--only-functions", help="Trace only these functions", nargs="+", default=None) @@ -654,16 +772,13 @@ def main(): "__cached__": None, } try: - tracer = Tracer( + Tracer( output=args.outfile, functions=args.only_functions, max_function_count=args.max_function_count, timeout=args.tracer_timeout, config_file_path=args.codeflash_config, - ) - - tracer.runctx(code, globs, None) - print(tracer.functions) + ).runctx(code, globs, None) except BrokenPipeError as exc: # Prevent "Exception ignored" during interpreter shutdown. diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 071535c6a..570888fcc 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,12 +1,27 @@ +from codeflash.benchmarking.codeflash_trace import codeflash_trace from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from pathlib import Path from codeflash.code_utils.code_utils import get_run_tmp_file +import shutil 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" - output_file = Path("test_trace_benchmarks.trace").resolve() - trace_benchmarks_pytest(benchmarks_root, project_root, output_file) + # make directory in project_root / "tests" + + + tests_root = project_root / "tests" / "test_trace_benchmarks" + tests_root.mkdir(parents=False, exist_ok=False) + output_file = (tests_root / Path("test_trace_benchmarks.trace")).resolve() + trace_benchmarks_pytest(benchmarks_root, tests_root, project_root, output_file) assert output_file.exists() - output_file.unlink() \ No newline at end of file + + test1_path = tests_root / Path("test_benchmark_bubble_sort_py_test_sort__replay_test_0.py") + assert test1_path.exists() + + # test1_code = """""" + # assert test1_path.read_text("utf-8").strip()==test1_code.strip() + # cleanup + # shutil.rmtree(tests_root) + # output_file.unlink() \ No newline at end of file From adffb9d501093e4cb88a7acc42fc324226b819c9 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 14 Mar 2025 15:03:36 -0700 Subject: [PATCH 062/122] replay test functionality working for functions, methods, static methods, class methods, init. basic instrumentation logic for codeflash_trace done. --- .../bubble_sort_codeflash_trace.py | 46 ++ .../benchmarks/test_benchmark_bubble_sort.py | 17 +- .../benchmarks/test_process_and_sort.py | 6 +- codeflash/benchmarking/codeflash_trace.py | 106 +---- .../instrument_codeflash_trace.py | 79 ++++ codeflash/benchmarking/plugin/plugin.py | 6 + .../pytest_new_process_trace_benchmarks.py | 2 +- codeflash/benchmarking/replay_test.py | 170 ++++++- tests/test_instrument_codeflash_capture.py | 441 ++++++------------ tests/test_instrument_codeflash_trace.py | 239 ++++++++++ tests/test_trace_benchmarks.py | 147 +++++- 11 files changed, 829 insertions(+), 430 deletions(-) create mode 100644 code_to_optimize/bubble_sort_codeflash_trace.py create mode 100644 tests/test_instrument_codeflash_trace.py diff --git a/code_to_optimize/bubble_sort_codeflash_trace.py b/code_to_optimize/bubble_sort_codeflash_trace.py new file mode 100644 index 000000000..ee4dbd999 --- /dev/null +++ b/code_to_optimize/bubble_sort_codeflash_trace.py @@ -0,0 +1,46 @@ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +@codeflash_trace +def sorter(arr): + for i in range(len(arr)): + for j in range(len(arr) - 1): + if arr[j] > arr[j + 1]: + temp = arr[j] + arr[j] = arr[j + 1] + arr[j + 1] = temp + return arr + +class Sorter: + @codeflash_trace + def __init__(self, arr): + self.arr = arr + @codeflash_trace + def sorter(self, multiplier): + for i in range(len(self.arr)): + for j in range(len(self.arr) - 1): + if self.arr[j] > self.arr[j + 1]: + temp = self.arr[j] + self.arr[j] = self.arr[j + 1] + self.arr[j + 1] = temp + return self.arr * multiplier + + @staticmethod + @codeflash_trace + def sort_static(arr): + for i in range(len(arr)): + for j in range(len(arr) - 1): + if arr[j] > arr[j + 1]: + temp = arr[j] + arr[j] = arr[j + 1] + arr[j + 1] = temp + return arr + + @classmethod + @codeflash_trace + def sort_class(cls, arr): + for i in range(len(arr)): + for j in range(len(arr) - 1): + if arr[j] > arr[j + 1]: + temp = arr[j] + arr[j] = arr[j + 1] + arr[j + 1] = temp + return arr diff --git a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py index 21c2bbb29..03b9d38d1 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py @@ -1,13 +1,20 @@ import pytest -from code_to_optimize.bubble_sort_codeflash_trace import sorter +from code_to_optimize.bubble_sort_codeflash_trace import sorter, Sorter def test_sort(benchmark): - result = benchmark(sorter, list(reversed(range(5000)))) - assert result == list(range(5000)) + result = benchmark(sorter, list(reversed(range(500)))) + assert result == list(range(500)) # This should not be picked up as a benchmark test def test_sort2(): - result = sorter(list(reversed(range(5000)))) - assert result == list(range(5000)) \ No newline at end of file + result = sorter(list(reversed(range(500)))) + assert result == list(range(500)) + +def test_class_sort(benchmark): + obj = Sorter(list(reversed(range(100)))) + result1 = benchmark(obj.sorter, 2) + result2 = benchmark(Sorter.sort_class, list(reversed(range(100)))) + result3 = benchmark(Sorter.sort_static, list(reversed(range(100)))) + result4 = benchmark(Sorter, [1,2,3]) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py index 2713721e4..bcd42eab9 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py @@ -1,8 +1,8 @@ from code_to_optimize.process_and_bubble_sort_codeflash_trace import compute_and_sort from code_to_optimize.bubble_sort_codeflash_trace import sorter def test_compute_and_sort(benchmark): - result = benchmark(compute_and_sort, list(reversed(range(5000)))) - assert result == 6247083.5 + result = benchmark(compute_and_sort, list(reversed(range(500)))) + assert result == 62208.5 def test_no_func(benchmark): - benchmark(sorter, list(reversed(range(5000)))) \ No newline at end of file + benchmark(sorter, list(reversed(range(500)))) \ No newline at end of file diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index c678b7643..65ba98783 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -51,6 +51,10 @@ def wrapper(*args, **kwargs): overhead_time = 0 try: + # Check if currently in pytest benchmark fixture + if os.environ.get("CODEFLASH_BENCHMARKING", "False") == "False": + return result + # Pickle the arguments pickled_args = pickle.dumps(args, protocol=pickle.HIGHEST_PROTOCOL) pickled_kwargs = pickle.dumps(kwargs, protocol=pickle.HIGHEST_PROTOCOL) @@ -58,6 +62,7 @@ def wrapper(*args, **kwargs): # 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", "") + benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "") # Calculate overhead time overhead_end_time = time.time() @@ -69,7 +74,7 @@ def wrapper(*args, **kwargs): class_name = qualname.split(".")[0] self.function_calls_data.append( (func.__name__, class_name, func.__module__, func.__code__.co_filename, - benchmark_function_name, benchmark_file_name, execution_time, + benchmark_function_name, benchmark_file_name, benchmark_line_number, execution_time, overhead_time, pickled_args, pickled_kwargs) ) @@ -100,7 +105,7 @@ def write_to_db(self, output_file: str) -> None: cur.execute( "CREATE TABLE IF NOT EXISTS function_calls(" "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," - "benchmark_function_name TEXT, benchmark_file_name TEXT, " + "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," "time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" ) @@ -108,8 +113,8 @@ def write_to_db(self, output_file: str) -> None: cur.executemany( "INSERT INTO function_calls " "(function_name, class_name, module_name, file_name, benchmark_function_name, " - "benchmark_file_name, time_ns, overhead_time_ns, args, kwargs) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + "benchmark_file_name, benchmark_line_number, time_ns, overhead_time_ns, args, kwargs) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", self.function_calls_data ) @@ -165,13 +170,14 @@ def print_codeflash_db(self, limit: int = None) -> None: print(f" File: {row[3]}") print(f" Benchmark Function: {row[4] or 'N/A'}") print(f" Benchmark File: {row[5] or 'N/A'}") - print(f" Execution Time: {row[6]:.6f} seconds") - print(f" Overhead Time: {row[7]:.6f} seconds") + print(f" Benchmark Line: {row[6] or 'N/A'}") + print(f" Execution Time: {row[7]:.6f} seconds") + print(f" Overhead Time: {row[8]:.6f} seconds") # Unpickle and print args and kwargs try: - args = pickle.loads(row[8]) - kwargs = pickle.loads(row[9]) + args = pickle.loads(row[9]) + kwargs = pickle.loads(row[10]) print(f" Args: {args}") print(f" Kwargs: {kwargs}") @@ -187,90 +193,6 @@ def print_codeflash_db(self, limit: int = None) -> None: except Exception as e: print(f"Error reading database: {e}") - def generate_replay_test(self, output_dir: str = None, project_root: str = "", test_framework: str = "pytest", - max_run_count: int = 100) -> None: - """ - Generate multiple replay tests from the traced function calls, grouping by benchmark name. - - Args: - output_dir: Directory to write the generated tests (if None, only returns the code) - project_root: Root directory of the project for module imports - test_framework: 'pytest' or 'unittest' - max_run_count: Maximum number of runs to include per function - - Returns: - Dictionary mapping benchmark names to generated test code - """ - import isort - from codeflash.verification.verification_utils import get_test_file_path - - if not self.db_path: - print("No database path set. Call write_to_db first or set db_path manually.") - return {} - - try: - # Import the function here to avoid circular imports - from codeflash.benchmarking.replay_test import create_trace_replay_test - - print("connecting to: ", self.db_path) - # Connect to the database - conn = sqlite3.connect(self.db_path) - cursor = conn.cursor() - - # Get distinct benchmark names - cursor.execute( - "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM function_calls" - ) - benchmarks = cursor.fetchall() - - # Generate a test for each benchmark - for benchmark in benchmarks: - benchmark_function_name, benchmark_file_name = benchmark - # Get functions associated with this benchmark - cursor.execute( - "SELECT DISTINCT function_name, class_name, module_name, file_name FROM function_calls " - "WHERE benchmark_function_name = ? AND benchmark_file_name = ?", - (benchmark_function_name, benchmark_file_name) - ) - - functions_data = [] - for func_row in cursor.fetchall(): - function_name, class_name, module_name, file_name = func_row - - # Add this function to our list - functions_data.append({ - "function_name": function_name, - "class_name": class_name, - "file_name": file_name, - "module_name": module_name - }) - - if not functions_data: - print(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") - continue - - # Generate the test code for this benchmark - test_code = create_trace_replay_test( - trace_file=self.db_path, - functions_data=functions_data, - test_framework=test_framework, - max_run_count=max_run_count, - ) - test_code = isort.code(test_code) - - # Write to file if requested - if output_dir: - output_file = get_test_file_path( - test_dir=Path(output_dir), function_name=f"{benchmark_file_name[5:]}_{benchmark_function_name}", test_type="replay" - ) - with open(output_file, 'w') as f: - f.write(test_code) - print(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") - - conn.close() - - except Exception as e: - print(f"Error generating replay tests: {e}") # Create a singleton instance codeflash_trace = CodeflashTrace() diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py index 8b1378917..99b2dad20 100644 --- a/codeflash/benchmarking/instrument_codeflash_trace.py +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -1 +1,80 @@ +import libcst as cst +from codeflash.discovery.functions_to_optimize import FunctionToOptimize + + +class AddDecoratorTransformer(cst.CSTTransformer): + def __init__(self, function_name, class_name=None): + super().__init__() + self.function_name = function_name + self.class_name = class_name + self.in_target_class = (class_name is None) # If no class name, always "in target class" + + def leave_ClassDef(self, original_node, updated_node): + if self.class_name and original_node.name.value == self.class_name: + self.in_target_class = False + return updated_node + + def visit_ClassDef(self, node): + if self.class_name and node.name.value == self.class_name: + self.in_target_class = True + return True + + def leave_FunctionDef(self, original_node, updated_node): + if not self.in_target_class or original_node.name.value != self.function_name: + return updated_node + + # Create the codeflash_trace decorator + decorator = cst.Decorator( + decorator=cst.Name(value="codeflash_trace") + ) + + # Add the new decorator after any existing decorators + updated_decorators = list(updated_node.decorators) + [decorator] + + # Return the updated node with the new decorator + return updated_node.with_changes( + decorators=updated_decorators + ) + + +def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToOptimize) -> str: + """Add codeflash_trace to a function. + + Args: + code: The source code as a string + function_to_optimize: The FunctionToOptimize instance containing function details + + Returns: + The modified source code as a string + """ + # Extract class name if present + class_name = None + if len(function_to_optimize.parents) == 1 and function_to_optimize.parents[0].type == "ClassDef": + class_name = function_to_optimize.parents[0].name + + transformer = AddDecoratorTransformer( + function_name=function_to_optimize.function_name, + class_name=class_name + ) + + module = cst.parse_module(code) + modified_module = module.visit(transformer) + return modified_module.code + + +def instrument_codeflash_trace( + function_to_optimize: FunctionToOptimize +) -> None: + """Instrument __init__ function with codeflash_trace decorator if it's in a class.""" + # Instrument fto class + original_code = function_to_optimize.file_path.read_text(encoding="utf-8") + + # Modify the code + modified_code = add_codeflash_decorator_to_code( + original_code, + function_to_optimize + ) + + # Write the modified code back to the file + function_to_optimize.file_path.write_text(modified_code, encoding="utf-8") diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index 6d8db9bf9..caf175a4e 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -1,3 +1,5 @@ +import sys + import pytest import time import os @@ -34,12 +36,16 @@ def benchmark(request): return None class Benchmark: + def __call__(self, func, *args, **kwargs): os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = request.node.name os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = request.node.fspath.basename + os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack + os.environ["CODEFLASH_BENCHMARKING"] = "True" start = time.process_time_ns() result = func(*args, **kwargs) end = time.process_time_ns() + os.environ["CODEFLASH_BENCHMARKING"] = "False" print(f"Benchmark: {func.__name__} took {end - start} ns") return result diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 8e1958fec..04c5e67ea 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -20,7 +20,7 @@ ) codeflash_trace.write_to_db(output_file) codeflash_trace.print_codeflash_db() - codeflash_trace.generate_replay_test(tests_root, project_root, test_framework="pytest") + except Exception as e: print(f"Failed to collect tests: {e!s}") exitcode = -1 \ No newline at end of file diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 9bc2c79f3..58ce456c2 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -5,6 +5,12 @@ from collections.abc import Generator from typing import Any, Dict +import isort + +from codeflash.cli_cmds.console import logger +from codeflash.discovery.functions_to_optimize import inspect_top_level_functions_or_methods +from codeflash.verification.verification_utils import get_test_file_path +from pathlib import Path def get_next_arg_and_return( trace_file: str, function_name: str, file_name: str, class_name: str | None = None, num_to_get: int = 25 @@ -20,21 +26,21 @@ def get_next_arg_and_return( ) else: cursor = cur.execute( - "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? ORDER BY time_ns ASC LIMIT ?", + "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? AND class_name = '' ORDER BY time_ns ASC LIMIT ?", (function_name, file_name, limit), ) while (val := cursor.fetchone()) is not None: - yield val[8], val[9] # args and kwargs are at indices 7 and 8 + yield val[9], val[10] # args and kwargs are at indices 7 and 8 def get_function_alias(module: str, function_name: str) -> str: return "_".join(module.split(".")) + "_" + function_name -def create_trace_replay_test( +def create_trace_replay_test_code( trace_file: str, - functions_data: list[Dict[str, Any]], + functions_data: list[dict[str, Any]], test_framework: str = "pytest", max_run_count=100 ) -> str: @@ -52,7 +58,7 @@ def create_trace_replay_test( """ assert test_framework in ["pytest", "unittest"] - imports = f"""import dill as pickle + imports = f"""import dill as pickle {"import unittest" if test_framework == "unittest" else ""} from codeflash.benchmarking.replay_test import get_next_arg_and_return """ @@ -62,7 +68,6 @@ def create_trace_replay_test( module_name = func.get("module_name") function_name = func.get("function_name") class_name = func.get("class_name", "") - if class_name: function_imports.append( f"from {module_name} import {class_name} as {get_function_alias(module_name, class_name)}" @@ -90,12 +95,37 @@ def create_trace_replay_test( """ ) + test_method_body = textwrap.dedent( + """\ + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl){filter_variables} + function_name = "{orig_function_name}" + if not args: + raise ValueError("No arguments provided for the method.") + if function_name == "__init__": + ret = {class_name_alias}(*args[1:], **kwargs) + else: + instance = args[0] # self + ret = instance{method_name}(*args[1:], **kwargs) + """) + test_class_method_body = textwrap.dedent( """\ for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} - ret = {class_name_alias}{method_name}(**args, **kwargs) + if not args: + raise ValueError("No arguments provided for the method.") + ret = {class_name_alias}{method_name}(*args[1:], **kwargs) + """ + ) + test_static_method_body = textwrap.dedent( + """\ + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl){filter_variables} + ret = {class_name_alias}{method_name}(*args, **kwargs) """ ) @@ -111,7 +141,9 @@ def create_trace_replay_test( function_name = func.get("function_name") class_name = func.get("class_name") file_name = func.get("file_name") - + function_properties = func.get("function_properties") + print(f"Class: {class_name}, Function: {function_name}") + print(function_properties) if not class_name: alias = get_function_alias(module_name, function_name) test_body = test_function_body.format( @@ -125,16 +157,38 @@ def create_trace_replay_test( alias = get_function_alias(module_name, class_name + "_" + function_name) filter_variables = "" + # filter_variables = '\n args.pop("cls", None)' method_name = "." + function_name if function_name != "__init__" else "" - test_body = test_class_method_body.format( - orig_function_name=function_name, - file_name=file_name, - class_name_alias=class_name_alias, - class_name=class_name, - method_name=method_name, - max_run_count=max_run_count, - filter_variables=filter_variables, - ) + if function_properties.is_classmethod: + test_body = test_class_method_body.format( + orig_function_name=function_name, + file_name=file_name, + class_name_alias=class_name_alias, + class_name=class_name, + method_name=method_name, + max_run_count=max_run_count, + filter_variables=filter_variables, + ) + elif function_properties.is_staticmethod: + test_body = test_static_method_body.format( + orig_function_name=function_name, + file_name=file_name, + class_name_alias=class_name_alias, + class_name=class_name, + method_name=method_name, + max_run_count=max_run_count, + filter_variables=filter_variables, + ) + else: + test_body = test_method_body.format( + orig_function_name=function_name, + file_name=file_name, + class_name_alias=class_name_alias, + class_name=class_name, + method_name=method_name, + max_run_count=max_run_count, + filter_variables=filter_variables, + ) formatted_test_body = textwrap.indent(test_body, " " if test_framework == "unittest" else " ") @@ -142,3 +196,85 @@ def create_trace_replay_test( test_template += f"def test_{alias}({self}):\n{formatted_test_body}\n" return imports + "\n" + metadata + "\n" + test_template + +def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework: str = "pytest", max_run_count: int = 100) -> None: + """Generate multiple replay tests from the traced function calls, grouping by benchmark name. + + Args: + trace_file_path: Path to the SQLite database file + output_dir: Directory to write the generated tests (if None, only returns the code) + project_root: Root directory of the project for module imports + test_framework: 'pytest' or 'unittest' + max_run_count: Maximum number of runs to include per function + + Returns: + Dictionary mapping benchmark names to generated test code + + """ + try: + # Connect to the database + conn = sqlite3.connect(trace_file_path.as_posix()) + cursor = conn.cursor() + + # Get distinct benchmark names + cursor.execute( + "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM function_calls" + ) + benchmarks = cursor.fetchall() + + # Generate a test for each benchmark + for benchmark in benchmarks: + benchmark_function_name, benchmark_file_name = benchmark + # Get functions associated with this benchmark + cursor.execute( + "SELECT DISTINCT function_name, class_name, module_name, file_name, benchmark_line_number FROM function_calls " + "WHERE benchmark_function_name = ? AND benchmark_file_name = ?", + (benchmark_function_name, benchmark_file_name) + ) + + functions_data = [] + for func_row in cursor.fetchall(): + function_name, class_name, module_name, file_name, benchmark_line_number = func_row + + # Add this function to our list + functions_data.append({ + "function_name": function_name, + "class_name": class_name, + "file_name": file_name, + "module_name": module_name, + "benchmark_function_name": benchmark_function_name, + "benchmark_file_name": benchmark_file_name, + "benchmark_line_number": benchmark_line_number, + "function_properties": inspect_top_level_functions_or_methods( + file_name=file_name, + function_or_method_name=function_name, + class_name=class_name, + ) + }) + + if not functions_data: + print(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") + continue + + # Generate the test code for this benchmark + test_code = create_trace_replay_test_code( + trace_file=trace_file_path.as_posix(), + functions_data=functions_data, + test_framework=test_framework, + max_run_count=max_run_count, + ) + test_code = isort.code(test_code) + + # Write to file if requested + if output_dir: + output_file = get_test_file_path( + test_dir=Path(output_dir), function_name=f"{benchmark_file_name[5:]}_{benchmark_function_name}", test_type="replay" + ) + with open(output_file, 'w') as f: + f.write(test_code) + print(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") + + conn.close() + + except Exception as e: + print(f"Error generating replay tests: {e}") diff --git a/tests/test_instrument_codeflash_capture.py b/tests/test_instrument_codeflash_capture.py index fe5a6bcd3..5cd5ce322 100644 --- a/tests/test_instrument_codeflash_capture.py +++ b/tests/test_instrument_codeflash_capture.py @@ -1,356 +1,193 @@ -from pathlib import Path +from __future__ import annotations -from codeflash.code_utils.code_utils import get_run_tmp_file -from codeflash.discovery.functions_to_optimize import FunctionToOptimize -from codeflash.models.models import FunctionParent -from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture +from codeflash.benchmarking.instrument_codeflash_trace import add_codeflash_decorator_to_code -def test_add_codeflash_capture(): - original_code = """ -class MyClass: - def __init__(self): - self.x = 1 - - def target_function(self): - return self.x + 1 +def test_add_decorator_to_normal_function() -> None: + """Test adding decorator to a normal function.""" + code = """ +def normal_function(): + return "Hello, World!" """ - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - expected = f""" -from codeflash.verification.codeflash_capture import codeflash_capture - - -class MyClass: - - @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) - def __init__(self): - self.x = 1 - def target_function(self): - return self.x + 1 -""" - test_path.write_text(original_code) - - function = FunctionToOptimize( - function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="normal_function" ) - try: - instrument_codeflash_capture(function, {}, test_path.parent) - modified_code = test_path.read_text() - assert modified_code.strip() == expected.strip() - - finally: - test_path.unlink(missing_ok=True) - - -def test_add_codeflash_capture_no_parent(): - original_code = """ -class MyClass: - - def target_function(self): - return self.x + 1 + expected_code = """ +@codeflash_trace +def normal_function(): + return "Hello, World!" """ - expected = """ -class MyClass: + assert modified_code.strip() == expected_code.strip() - def target_function(self): - return self.x + 1 +def test_add_decorator_to_normal_method() -> None: + """Test adding decorator to a normal method.""" + code = """ +class TestClass: + def normal_method(self): + return "Hello from method" """ - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - test_path.write_text(original_code) - - function = FunctionToOptimize(function_name="target_function", file_path=test_path, parents=[]) - try: - instrument_codeflash_capture(function, {}, test_path.parent) - modified_code = test_path.read_text() - assert modified_code.strip() == expected.strip() - finally: - test_path.unlink(missing_ok=True) - - -def test_add_codeflash_capture_no_init(): - # Test input code - original_code = """ -class MyClass(ParentClass): + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="normal_method", + class_name="TestClass" + ) - def target_function(self): - return self.x + 1 + expected_code = """ +class TestClass: + @codeflash_trace + def normal_method(self): + return "Hello from method" """ - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - expected = f""" -from codeflash.verification.codeflash_capture import codeflash_capture + assert modified_code.strip() == expected_code.strip() -class MyClass(ParentClass): - - @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - - def target_function(self): - return self.x + 1 +def test_add_decorator_to_classmethod() -> None: + """Test adding decorator to a classmethod.""" + code = """ +class TestClass: + @classmethod + def class_method(cls): + return "Hello from classmethod" """ - test_path.write_text(original_code) - function = FunctionToOptimize( - function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="class_method", + class_name="TestClass" ) - try: - instrument_codeflash_capture(function, {}, test_path.parent) - modified_code = test_path.read_text() - assert modified_code.strip() == expected.strip() - - finally: - test_path.unlink(missing_ok=True) - - -def test_add_codeflash_capture_with_helpers(): - # Test input code - original_code = """ -class MyClass: - def __init__(self): - self.x = 1 - - def target_function(self): - return helper() + 1 - - def helper(self): - return self.x + expected_code = """ +class TestClass: + @classmethod + @codeflash_trace + def class_method(cls): + return "Hello from classmethod" """ - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - expected = f""" -from codeflash.verification.codeflash_capture import codeflash_capture + assert modified_code.strip() == expected_code.strip() -class MyClass: - - @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) - def __init__(self): - self.x = 1 - - def target_function(self): - return helper() + 1 - - def helper(self): - return self.x +def test_add_decorator_to_staticmethod() -> None: + """Test adding decorator to a staticmethod.""" + code = """ +class TestClass: + @staticmethod + def static_method(): + return "Hello from staticmethod" """ - test_path.write_text(original_code) - - function = FunctionToOptimize( - function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="static_method", + class_name="TestClass" ) - try: - instrument_codeflash_capture( - function, {test_path: {"MyClass"}}, test_path.parent - ) # MyClass was removed from the file_path_to_helper_class as it shares class with FTO - modified_code = test_path.read_text() - assert modified_code.strip() == expected.strip() - - finally: - test_path.unlink(missing_ok=True) - - -def test_add_codeflash_capture_with_helpers_2(): - # Test input code - original_code = """ -from test_helper_file import HelperClass - -class MyClass: - def __init__(self): - self.x = 1 - - def target_function(self): - return HelperClass().helper() + 1 -""" - original_helper = """ -class HelperClass: - def __init__(self): - self.y = 1 - def helper(self): - return 1 -""" - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - expected = f""" -from test_helper_file import HelperClass - -from codeflash.verification.codeflash_capture import codeflash_capture - - -class MyClass: - - @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) - def __init__(self): - self.x = 1 - - def target_function(self): - return HelperClass().helper() + 1 + expected_code = """ +class TestClass: + @staticmethod + @codeflash_trace + def static_method(): + return "Hello from staticmethod" """ - expected_helper = f""" -from codeflash.verification.codeflash_capture import codeflash_capture - -class HelperClass: + assert modified_code.strip() == expected_code.strip() - @codeflash_capture(function_name='HelperClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) - def __init__(self): - self.y = 1 - - def helper(self): - return 1 +def test_add_decorator_to_init_function() -> None: + """Test adding decorator to an __init__ function.""" + code = """ +class TestClass: + def __init__(self, value): + self.value = value """ - test_path.write_text(original_code) - helper_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_helper_file.py").resolve() - helper_path.write_text(original_helper) - - function = FunctionToOptimize( - function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="__init__", + class_name="TestClass" ) - try: - instrument_codeflash_capture(function, {helper_path: {"HelperClass"}}, test_path.parent) - modified_code = test_path.read_text() - assert modified_code.strip() == expected.strip() - assert helper_path.read_text().strip() == expected_helper.strip() - finally: - test_path.unlink(missing_ok=True) - helper_path.unlink(missing_ok=True) - - -def test_add_codeflash_capture_with_multiple_helpers(): - # Test input code with imports from two helper files - original_code = """ -from helper_file_1 import HelperClass1 -from helper_file_2 import HelperClass2, AnotherHelperClass - -class MyClass: - def __init__(self): - self.x = 1 - - def target_function(self): - helper1 = HelperClass1().helper1() - helper2 = HelperClass2().helper2() - another = AnotherHelperClass().another_helper() - return helper1 + helper2 + another + expected_code = """ +class TestClass: + @codeflash_trace + def __init__(self, value): + self.value = value """ - # First helper file content - original_helper1 = """ -class HelperClass1: - def __init__(self): - self.y = 1 - def helper1(self): - return 1 -""" + assert modified_code.strip() == expected_code.strip() - # Second helper file content - original_helper2 = """ -class HelperClass2: - def __init__(self): - self.z = 2 - def helper2(self): - return 2 - -class AnotherHelperClass: - def another_helper(self): - return 3 +def test_add_decorator_with_multiple_decorators() -> None: + """Test adding decorator to a function with multiple existing decorators.""" + code = """ +class TestClass: + @property + @other_decorator + def property_method(self): + return self._value """ - test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() - expected = f""" -from helper_file_1 import HelperClass1 -from helper_file_2 import AnotherHelperClass, HelperClass2 -from codeflash.verification.codeflash_capture import codeflash_capture - - -class MyClass: - - @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) - def __init__(self): - self.x = 1 + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="property_method", + class_name="TestClass" + ) - def target_function(self): - helper1 = HelperClass1().helper1() - helper2 = HelperClass2().helper2() - another = AnotherHelperClass().another_helper() - return helper1 + helper2 + another + expected_code = """ +class TestClass: + @property + @other_decorator + @codeflash_trace + def property_method(self): + return self._value """ - # Expected output for first helper file - expected_helper1 = f""" -from codeflash.verification.codeflash_capture import codeflash_capture + assert modified_code.strip() == expected_code.strip() +def test_add_decorator_to_function_in_multiple_classes() -> None: + """Test that only the right class's method gets the decorator.""" + code = """ +class TestClass: + def test_method(self): + return "This should get decorated" -class HelperClass1: - - @codeflash_capture(function_name='HelperClass1.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) - def __init__(self): - self.y = 1 - - def helper1(self): - return 1 +class OtherClass: + def test_method(self): + return "This should NOT get decorated" """ - # Expected output for second helper file - expected_helper2 = f""" -from codeflash.verification.codeflash_capture import codeflash_capture - - -class HelperClass2: - - @codeflash_capture(function_name='HelperClass2.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) - def __init__(self): - self.z = 2 - - def helper2(self): - return 2 - -class AnotherHelperClass: + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="test_method", + class_name="TestClass" + ) - @codeflash_capture(function_name='AnotherHelperClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + expected_code = """ +class TestClass: + @codeflash_trace + def test_method(self): + return "This should get decorated" - def another_helper(self): - return 3 +class OtherClass: + def test_method(self): + return "This should NOT get decorated" """ - # Set up test files - helper1_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/helper_file_1.py").resolve() - helper2_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/helper_file_2.py").resolve() + assert modified_code.strip() == expected_code.strip() - # Write original content to files - test_path.write_text(original_code) - helper1_path.write_text(original_helper1) - helper2_path.write_text(original_helper2) +def test_add_decorator_to_nonexistent_function() -> None: + """Test that code remains unchanged when function doesn't exist.""" + code = """ +def existing_function(): + return "This exists" +""" - # Create FunctionToOptimize instance - function = FunctionToOptimize( - function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] + modified_code = add_codeflash_decorator_to_code( + code=code, + function_name="nonexistent_function" ) - try: - # Instrument code with multiple helper files - helper_classes = {helper1_path: {"HelperClass1"}, helper2_path: {"HelperClass2", "AnotherHelperClass"}} - instrument_codeflash_capture(function, helper_classes, test_path.parent) - - # Verify the modifications - modified_code = test_path.read_text() - modified_helper1 = helper1_path.read_text() - modified_helper2 = helper2_path.read_text() - - assert modified_code.strip() == expected.strip() - assert modified_helper1.strip() == expected_helper1.strip() - assert modified_helper2.strip() == expected_helper2.strip() - - finally: - # Clean up test files - test_path.unlink(missing_ok=True) - helper1_path.unlink(missing_ok=True) - helper2_path.unlink(missing_ok=True) + # Code should remain unchanged + assert modified_code.strip() == code.strip() diff --git a/tests/test_instrument_codeflash_trace.py b/tests/test_instrument_codeflash_trace.py new file mode 100644 index 000000000..56008faa9 --- /dev/null +++ b/tests/test_instrument_codeflash_trace.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +from pathlib import Path + +from codeflash.benchmarking.instrument_codeflash_trace import add_codeflash_decorator_to_code + +from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize + + +def test_add_decorator_to_normal_function() -> None: + """Test adding decorator to a normal function.""" + code = """ +def normal_function(): + return "Hello, World!" +""" + + fto = FunctionToOptimize( + function_name="normal_function", + file_path=Path("dummy_path.py"), + parents=[] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +@codeflash_trace +def normal_function(): + return "Hello, World!" +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_normal_method() -> None: + """Test adding decorator to a normal method.""" + code = """ +class TestClass: + def normal_method(self): + return "Hello from method" +""" + + fto = FunctionToOptimize( + function_name="normal_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @codeflash_trace + def normal_method(self): + return "Hello from method" +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_classmethod() -> None: + """Test adding decorator to a classmethod.""" + code = """ +class TestClass: + @classmethod + def class_method(cls): + return "Hello from classmethod" +""" + + fto = FunctionToOptimize( + function_name="class_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @classmethod + @codeflash_trace + def class_method(cls): + return "Hello from classmethod" +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_staticmethod() -> None: + """Test adding decorator to a staticmethod.""" + code = """ +class TestClass: + @staticmethod + def static_method(): + return "Hello from staticmethod" +""" + + fto = FunctionToOptimize( + function_name="static_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @staticmethod + @codeflash_trace + def static_method(): + return "Hello from staticmethod" +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_init_function() -> None: + """Test adding decorator to an __init__ function.""" + code = """ +class TestClass: + def __init__(self, value): + self.value = value +""" + + fto = FunctionToOptimize( + function_name="__init__", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @codeflash_trace + def __init__(self, value): + self.value = value +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_with_multiple_decorators() -> None: + """Test adding decorator to a function with multiple existing decorators.""" + code = """ +class TestClass: + @property + @other_decorator + def property_method(self): + return self._value +""" + + fto = FunctionToOptimize( + function_name="property_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @property + @other_decorator + @codeflash_trace + def property_method(self): + return self._value +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_function_in_multiple_classes() -> None: + """Test that only the right class's method gets the decorator.""" + code = """ +class TestClass: + def test_method(self): + return "This should get decorated" + +class OtherClass: + def test_method(self): + return "This should NOT get decorated" +""" + + fto = FunctionToOptimize( + function_name="test_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + expected_code = """ +class TestClass: + @codeflash_trace + def test_method(self): + return "This should get decorated" + +class OtherClass: + def test_method(self): + return "This should NOT get decorated" +""" + + assert modified_code.strip() == expected_code.strip() + +def test_add_decorator_to_nonexistent_function() -> None: + """Test that code remains unchanged when function doesn't exist.""" + code = """ +def existing_function(): + return "This exists" +""" + + fto = FunctionToOptimize( + function_name="nonexistent_function", + file_path=Path("dummy_path.py"), + parents=[] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + function_to_optimize=fto + ) + + # Code should remain unchanged + assert modified_code.strip() == code.strip() diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 570888fcc..c49e7c693 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,27 +1,154 @@ +import sqlite3 + from codeflash.benchmarking.codeflash_trace import codeflash_trace from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest +from codeflash.benchmarking.replay_test import generate_replay_test from pathlib import Path from codeflash.code_utils.code_utils import get_run_tmp_file import shutil + 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" - # make directory in project_root / "tests" - - tests_root = project_root / "tests" / "test_trace_benchmarks" tests_root.mkdir(parents=False, exist_ok=False) output_file = (tests_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_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM function_calls ORDER BY benchmark_file_name, benchmark_function_name, function_name") + 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)}" + + # Expected function calls + expected_calls = [ + ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_class_sort", "test_benchmark_bubble_sort.py", 20), + + ("sort_class", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_class_sort", "test_benchmark_bubble_sort.py", 18), + + ("sort_static", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_class_sort", "test_benchmark_bubble_sort.py", 19), + + ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_class_sort", "test_benchmark_bubble_sort.py", 17), + + ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_sort", "test_benchmark_bubble_sort.py", 7), + + ("compute_and_sort", "", "code_to_optimize.process_and_bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/process_and_bubble_sort_codeflash_trace.py'}", + "test_compute_and_sort", "test_process_and_sort.py", 4), + + ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", + f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + "test_no_func", "test_process_and_sort.py", 8), + ] + 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_name" + 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_file_name" + assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number" + # Close connection + conn.close() + generate_replay_test(output_file, tests_root) + test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort_py_test_class_sort__replay_test_0.py") + assert test_class_sort_path.exists() + test_class_sort_code = f""" +import dill as pickle + +from code_to_optimize.bubble_sort_codeflash_trace import \\ + Sorter as code_to_optimize_bubble_sort_codeflash_trace_Sorter +from codeflash.benchmarking.replay_test import get_next_arg_and_return + +functions = ['sorter', 'sort_class', 'sort_static'] +trace_file_path = r"{output_file.as_posix()}" + +def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + function_name = "sorter" + if not args: + raise ValueError("No arguments provided for the method.") + if function_name == "__init__": + ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter(*args[1:], **kwargs) + else: + instance = args[0] # self + ret = instance.sorter(*args[1:], **kwargs) + +def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + if not args: + raise ValueError("No arguments provided for the method.") + ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_class(*args[1:], **kwargs) + +def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_static(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_static(*args, **kwargs) + +def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + function_name = "__init__" + if not args: + raise ValueError("No arguments provided for the method.") + if function_name == "__init__": + ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter(*args[1:], **kwargs) + else: + instance = args[0] # self + ret = instance(*args[1:], **kwargs) + +""" + assert test_class_sort_path.read_text("utf-8").strip()==test_class_sort_code.strip() + + test_sort_path = tests_root / Path("test_benchmark_bubble_sort_py_test_sort__replay_test_0.py") + assert test_sort_path.exists() + test_sort_code = f""" +import dill as pickle + +from code_to_optimize.bubble_sort_codeflash_trace import \\ + sorter as code_to_optimize_bubble_sort_codeflash_trace_sorter +from codeflash.benchmarking.replay_test import get_next_arg_and_return + +functions = ['sorter'] +trace_file_path = r"{output_file}" - test1_path = tests_root / Path("test_benchmark_bubble_sort_py_test_sort__replay_test_0.py") - assert test1_path.exists() +def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + ret = code_to_optimize_bubble_sort_codeflash_trace_sorter(*args, **kwargs) - # test1_code = """""" - # assert test1_path.read_text("utf-8").strip()==test1_code.strip() - # cleanup - # shutil.rmtree(tests_root) - # output_file.unlink() \ No newline at end of file +""" + assert test_sort_path.read_text("utf-8").strip()==test_sort_code.strip() + finally: + # cleanup + shutil.rmtree(tests_root) + pass \ No newline at end of file From f9144ec83988c45cf501b443342818fb89a81ea3 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 14 Mar 2025 15:05:55 -0700 Subject: [PATCH 063/122] restored overwritten logic --- tests/test_instrument_codeflash_capture.py | 441 ++++++++++++++------- 1 file changed, 302 insertions(+), 139 deletions(-) diff --git a/tests/test_instrument_codeflash_capture.py b/tests/test_instrument_codeflash_capture.py index 5cd5ce322..fe5a6bcd3 100644 --- a/tests/test_instrument_codeflash_capture.py +++ b/tests/test_instrument_codeflash_capture.py @@ -1,193 +1,356 @@ -from __future__ import annotations +from pathlib import Path -from codeflash.benchmarking.instrument_codeflash_trace import add_codeflash_decorator_to_code +from codeflash.code_utils.code_utils import get_run_tmp_file +from codeflash.discovery.functions_to_optimize import FunctionToOptimize +from codeflash.models.models import FunctionParent +from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture -def test_add_decorator_to_normal_function() -> None: - """Test adding decorator to a normal function.""" - code = """ -def normal_function(): - return "Hello, World!" +def test_add_codeflash_capture(): + original_code = """ +class MyClass: + def __init__(self): + self.x = 1 + + def target_function(self): + return self.x + 1 """ + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + expected = f""" +from codeflash.verification.codeflash_capture import codeflash_capture - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="normal_function" - ) - expected_code = """ -@codeflash_trace -def normal_function(): - return "Hello, World!" -""" +class MyClass: - assert modified_code.strip() == expected_code.strip() + @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) + def __init__(self): + self.x = 1 -def test_add_decorator_to_normal_method() -> None: - """Test adding decorator to a normal method.""" - code = """ -class TestClass: - def normal_method(self): - return "Hello from method" + def target_function(self): + return self.x + 1 """ + test_path.write_text(original_code) - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="normal_method", - class_name="TestClass" + function = FunctionToOptimize( + function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] ) - expected_code = """ -class TestClass: - @codeflash_trace - def normal_method(self): - return "Hello from method" + try: + instrument_codeflash_capture(function, {}, test_path.parent) + modified_code = test_path.read_text() + assert modified_code.strip() == expected.strip() + + finally: + test_path.unlink(missing_ok=True) + + +def test_add_codeflash_capture_no_parent(): + original_code = """ +class MyClass: + + def target_function(self): + return self.x + 1 """ - assert modified_code.strip() == expected_code.strip() + expected = """ +class MyClass: -def test_add_decorator_to_classmethod() -> None: - """Test adding decorator to a classmethod.""" - code = """ -class TestClass: - @classmethod - def class_method(cls): - return "Hello from classmethod" + def target_function(self): + return self.x + 1 """ + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + test_path.write_text(original_code) - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="class_method", - class_name="TestClass" - ) + function = FunctionToOptimize(function_name="target_function", file_path=test_path, parents=[]) + + try: + instrument_codeflash_capture(function, {}, test_path.parent) + modified_code = test_path.read_text() + assert modified_code.strip() == expected.strip() + finally: + test_path.unlink(missing_ok=True) + + +def test_add_codeflash_capture_no_init(): + # Test input code + original_code = """ +class MyClass(ParentClass): - expected_code = """ -class TestClass: - @classmethod - @codeflash_trace - def class_method(cls): - return "Hello from classmethod" + def target_function(self): + return self.x + 1 """ + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + expected = f""" +from codeflash.verification.codeflash_capture import codeflash_capture - assert modified_code.strip() == expected_code.strip() -def test_add_decorator_to_staticmethod() -> None: - """Test adding decorator to a staticmethod.""" - code = """ -class TestClass: - @staticmethod - def static_method(): - return "Hello from staticmethod" +class MyClass(ParentClass): + + @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def target_function(self): + return self.x + 1 """ + test_path.write_text(original_code) - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="static_method", - class_name="TestClass" + function = FunctionToOptimize( + function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] ) - expected_code = """ -class TestClass: - @staticmethod - @codeflash_trace - def static_method(): - return "Hello from staticmethod" + try: + instrument_codeflash_capture(function, {}, test_path.parent) + modified_code = test_path.read_text() + assert modified_code.strip() == expected.strip() + + finally: + test_path.unlink(missing_ok=True) + + +def test_add_codeflash_capture_with_helpers(): + # Test input code + original_code = """ +class MyClass: + def __init__(self): + self.x = 1 + + def target_function(self): + return helper() + 1 + + def helper(self): + return self.x """ + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + expected = f""" +from codeflash.verification.codeflash_capture import codeflash_capture + + +class MyClass: + + @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) + def __init__(self): + self.x = 1 - assert modified_code.strip() == expected_code.strip() + def target_function(self): + return helper() + 1 -def test_add_decorator_to_init_function() -> None: - """Test adding decorator to an __init__ function.""" - code = """ -class TestClass: - def __init__(self, value): - self.value = value + def helper(self): + return self.x """ - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="__init__", - class_name="TestClass" + test_path.write_text(original_code) + + function = FunctionToOptimize( + function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] ) - expected_code = """ -class TestClass: - @codeflash_trace - def __init__(self, value): - self.value = value + try: + instrument_codeflash_capture( + function, {test_path: {"MyClass"}}, test_path.parent + ) # MyClass was removed from the file_path_to_helper_class as it shares class with FTO + modified_code = test_path.read_text() + assert modified_code.strip() == expected.strip() + + finally: + test_path.unlink(missing_ok=True) + + +def test_add_codeflash_capture_with_helpers_2(): + # Test input code + original_code = """ +from test_helper_file import HelperClass + +class MyClass: + def __init__(self): + self.x = 1 + + def target_function(self): + return HelperClass().helper() + 1 """ + original_helper = """ +class HelperClass: + def __init__(self): + self.y = 1 + def helper(self): + return 1 +""" + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + expected = f""" +from test_helper_file import HelperClass + +from codeflash.verification.codeflash_capture import codeflash_capture + + +class MyClass: + + @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) + def __init__(self): + self.x = 1 + + def target_function(self): + return HelperClass().helper() + 1 +""" + expected_helper = f""" +from codeflash.verification.codeflash_capture import codeflash_capture - assert modified_code.strip() == expected_code.strip() -def test_add_decorator_with_multiple_decorators() -> None: - """Test adding decorator to a function with multiple existing decorators.""" - code = """ -class TestClass: - @property - @other_decorator - def property_method(self): - return self._value +class HelperClass: + + @codeflash_capture(function_name='HelperClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) + def __init__(self): + self.y = 1 + + def helper(self): + return 1 """ - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="property_method", - class_name="TestClass" + test_path.write_text(original_code) + helper_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_helper_file.py").resolve() + helper_path.write_text(original_helper) + + function = FunctionToOptimize( + function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] ) - expected_code = """ -class TestClass: - @property - @other_decorator - @codeflash_trace - def property_method(self): - return self._value + try: + instrument_codeflash_capture(function, {helper_path: {"HelperClass"}}, test_path.parent) + modified_code = test_path.read_text() + assert modified_code.strip() == expected.strip() + assert helper_path.read_text().strip() == expected_helper.strip() + finally: + test_path.unlink(missing_ok=True) + helper_path.unlink(missing_ok=True) + + +def test_add_codeflash_capture_with_multiple_helpers(): + # Test input code with imports from two helper files + original_code = """ +from helper_file_1 import HelperClass1 +from helper_file_2 import HelperClass2, AnotherHelperClass + +class MyClass: + def __init__(self): + self.x = 1 + + def target_function(self): + helper1 = HelperClass1().helper1() + helper2 = HelperClass2().helper2() + another = AnotherHelperClass().another_helper() + return helper1 + helper2 + another +""" + + # First helper file content + original_helper1 = """ +class HelperClass1: + def __init__(self): + self.y = 1 + def helper1(self): + return 1 +""" + + # Second helper file content + original_helper2 = """ +class HelperClass2: + def __init__(self): + self.z = 2 + def helper2(self): + return 2 + +class AnotherHelperClass: + def another_helper(self): + return 3 """ + test_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_file.py").resolve() + expected = f""" +from helper_file_1 import HelperClass1 +from helper_file_2 import AnotherHelperClass, HelperClass2 - assert modified_code.strip() == expected_code.strip() +from codeflash.verification.codeflash_capture import codeflash_capture -def test_add_decorator_to_function_in_multiple_classes() -> None: - """Test that only the right class's method gets the decorator.""" - code = """ -class TestClass: - def test_method(self): - return "This should get decorated" -class OtherClass: - def test_method(self): - return "This should NOT get decorated" +class MyClass: + + @codeflash_capture(function_name='MyClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=True) + def __init__(self): + self.x = 1 + + def target_function(self): + helper1 = HelperClass1().helper1() + helper2 = HelperClass2().helper2() + another = AnotherHelperClass().another_helper() + return helper1 + helper2 + another """ - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="test_method", - class_name="TestClass" - ) + # Expected output for first helper file + expected_helper1 = f""" +from codeflash.verification.codeflash_capture import codeflash_capture + - expected_code = """ -class TestClass: - @codeflash_trace - def test_method(self): - return "This should get decorated" +class HelperClass1: -class OtherClass: - def test_method(self): - return "This should NOT get decorated" + @codeflash_capture(function_name='HelperClass1.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) + def __init__(self): + self.y = 1 + + def helper1(self): + return 1 """ - assert modified_code.strip() == expected_code.strip() + # Expected output for second helper file + expected_helper2 = f""" +from codeflash.verification.codeflash_capture import codeflash_capture + + +class HelperClass2: -def test_add_decorator_to_nonexistent_function() -> None: - """Test that code remains unchanged when function doesn't exist.""" - code = """ -def existing_function(): - return "This exists" + @codeflash_capture(function_name='HelperClass2.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) + def __init__(self): + self.z = 2 + + def helper2(self): + return 2 + +class AnotherHelperClass: + + @codeflash_capture(function_name='AnotherHelperClass.__init__', tmp_dir_path='{get_run_tmp_file(Path("test_return_values"))!s}', tests_root='{test_path.parent!s}', is_fto=False) + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def another_helper(self): + return 3 """ - modified_code = add_codeflash_decorator_to_code( - code=code, - function_name="nonexistent_function" + # Set up test files + helper1_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/helper_file_1.py").resolve() + helper2_path = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/helper_file_2.py").resolve() + + # Write original content to files + test_path.write_text(original_code) + helper1_path.write_text(original_helper1) + helper2_path.write_text(original_helper2) + + # Create FunctionToOptimize instance + function = FunctionToOptimize( + function_name="target_function", file_path=test_path, parents=[FunctionParent(type="ClassDef", name="MyClass")] ) - # Code should remain unchanged - assert modified_code.strip() == code.strip() + try: + # Instrument code with multiple helper files + helper_classes = {helper1_path: {"HelperClass1"}, helper2_path: {"HelperClass2", "AnotherHelperClass"}} + instrument_codeflash_capture(function, helper_classes, test_path.parent) + + # Verify the modifications + modified_code = test_path.read_text() + modified_helper1 = helper1_path.read_text() + modified_helper2 = helper2_path.read_text() + + assert modified_code.strip() == expected.strip() + assert modified_helper1.strip() == expected_helper1.strip() + assert modified_helper2.strip() == expected_helper2.strip() + + finally: + # Clean up test files + test_path.unlink(missing_ok=True) + helper1_path.unlink(missing_ok=True) + helper2_path.unlink(missing_ok=True) From c29c8bffbac44f97bd78f2c247ee2c61553a5c9f Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 18 Mar 2025 10:00:05 -0700 Subject: [PATCH 064/122] functioning end to end, gets the funciton impact on benchmarks --- .../benchmarking/benchmark_database_utils.py | 179 ++++++++++++++++++ codeflash/benchmarking/codeflash_trace.py | 124 +----------- codeflash/benchmarking/get_trace_info.py | 165 ++++++++-------- .../instrument_codeflash_trace.py | 37 +++- codeflash/benchmarking/plugin/plugin.py | 20 +- .../pytest_new_process_trace_benchmarks.py | 16 +- codeflash/benchmarking/replay_test.py | 5 +- codeflash/benchmarking/trace_benchmarks.py | 6 +- codeflash/discovery/functions_to_optimize.py | 43 +++-- .../pytest_new_process_discover_benchmarks.py | 54 ------ codeflash/optimization/function_optimizer.py | 17 +- codeflash/optimization/optimizer.py | 25 ++- tests/test_instrument_codeflash_trace.py | 7 + 13 files changed, 394 insertions(+), 304 deletions(-) create mode 100644 codeflash/benchmarking/benchmark_database_utils.py delete mode 100644 codeflash/discovery/pytest_new_process_discover_benchmarks.py diff --git a/codeflash/benchmarking/benchmark_database_utils.py b/codeflash/benchmarking/benchmark_database_utils.py new file mode 100644 index 000000000..b9b36079d --- /dev/null +++ b/codeflash/benchmarking/benchmark_database_utils.py @@ -0,0 +1,179 @@ +import sqlite3 +from pathlib import Path + +import pickle + + +class BenchmarkDatabaseUtils: + def __init__(self, trace_path :Path) -> None: + self.trace_path = trace_path + self.connection = None + + def setup(self) -> None: + try: + # Open connection + self.connection = sqlite3.connect(self.trace_path) + cur = self.connection.cursor() + cur.execute("PRAGMA synchronous = OFF") + cur.execute( + "CREATE TABLE IF NOT EXISTS function_calls(" + "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," + "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," + "time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" + ) + cur.execute( + "CREATE TABLE IF NOT EXISTS benchmark_timings(" + "benchmark_file_name TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER," + "time_ns INTEGER)" # Added closing parenthesis + ) + self.connection.commit() + # Don't close the connection here + except Exception as e: + print(f"Database setup error: {e}") + if self.connection: + self.connection.close() + self.connection = None + raise + + def write_function_timings(self, data: list[tuple]) -> None: + if not self.connection: + self.connection = sqlite3.connect(self.trace_path) + + try: + cur = self.connection.cursor() + # Insert data into the function_calls table + cur.executemany( + "INSERT INTO function_calls " + "(function_name, class_name, module_name, file_name, benchmark_function_name, " + "benchmark_file_name, benchmark_line_number, time_ns, overhead_time_ns, args, kwargs) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + data + ) + self.connection.commit() + except Exception as e: + print(f"Error writing to function timings database: {e}") + self.connection.rollback() + raise + + def write_benchmark_timings(self, data: list[tuple]) -> None: + if not self.connection: + self.connection = sqlite3.connect(self.trace_path) + + try: + cur = self.connection.cursor() + # Insert data into the benchmark_timings table + cur.executemany( + "INSERT INTO benchmark_timings (benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns) VALUES (?, ?, ?, ?)", + data + ) + self.connection.commit() + except Exception as e: + print(f"Error writing to benchmark timings database: {e}") + self.connection.rollback() + raise + + def print_function_timings(self, limit: int = None) -> None: + """Print the contents of a CodeflashTrace SQLite database. + + Args: + limit: Maximum number of records to print (None for all) + """ + if not self.connection: + self.connection = sqlite3.connect(self.trace_path) + try: + cur = self.connection.cursor() + + # Get the count of records + cur.execute("SELECT COUNT(*) FROM function_calls") + total_records = cur.fetchone()[0] + print(f"Found {total_records} function call records in {self.trace_path}") + + # Build the query with optional limit + query = "SELECT * FROM function_calls" + if limit: + query += f" LIMIT {limit}" + + # Execute the query + cur.execute(query) + + # Print column names + columns = [desc[0] for desc in cur.description] + print("\nColumns:", columns) + print("\n" + "=" * 80 + "\n") + + # Print each row + for i, row in enumerate(cur.fetchall()): + print(f"Record #{i + 1}:") + print(f" Function: {row[0]}") + print(f" Class: {row[1]}") + print(f" Module: {row[2]}") + print(f" File: {row[3]}") + print(f" Benchmark Function: {row[4] or 'N/A'}") + print(f" Benchmark File: {row[5] or 'N/A'}") + print(f" Benchmark Line: {row[6] or 'N/A'}") + print(f" Execution Time: {row[7]:.6f} seconds") + print(f" Overhead Time: {row[8]:.6f} seconds") + + # Unpickle and print args and kwargs + try: + args = pickle.loads(row[9]) + kwargs = pickle.loads(row[10]) + + print(f" Args: {args}") + print(f" Kwargs: {kwargs}") + except Exception as e: + print(f" Error unpickling args/kwargs: {e}") + print(f" Raw args: {row[9]}") + print(f" Raw kwargs: {row[10]}") + + print("\n" + "-" * 40 + "\n") + + except Exception as e: + print(f"Error reading database: {e}") + + def print_benchmark_timings(self, limit: int = None) -> None: + """Print the contents of a CodeflashTrace SQLite database. + Args: + limit: Maximum number of records to print (None for all) + """ + if not self.connection: + self.connection = sqlite3.connect(self.trace_path) + try: + cur = self.connection.cursor() + + # Get the count of records + cur.execute("SELECT COUNT(*) FROM benchmark_timings") + total_records = cur.fetchone()[0] + print(f"Found {total_records} benchmark timing records in {self.trace_path}") + + # Build the query with optional limit + query = "SELECT * FROM benchmark_timings" + if limit: + query += f" LIMIT {limit}" + + # Execute the query + cur.execute(query) + + # Print column names + columns = [desc[0] for desc in cur.description] + print("\nColumns:", columns) + print("\n" + "=" * 80 + "\n") + + # Print each row + for i, row in enumerate(cur.fetchall()): + print(f"Record #{i + 1}:") + print(f" Benchmark File: {row[0] or 'N/A'}") + print(f" Benchmark Function: {row[1] or 'N/A'}") + print(f" Benchmark Line: {row[2] or 'N/A'}") + print(f" Execution Time: {row[3] / 1e9:.6f} seconds") # Convert nanoseconds to seconds + print("\n" + "-" * 40 + "\n") + + except Exception as e: + print(f"Error reading benchmark timings database: {e}") + + + def close(self) -> None: + if self.connection: + self.connection.close() + self.connection = None + diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 65ba98783..9b9afead7 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -39,16 +39,15 @@ def __call__(self, func: Callable) -> Callable: @functools.wraps(func) def wrapper(*args, **kwargs): # Measure execution time - start_time = time.time() + start_time = time.perf_counter_ns() result = func(*args, **kwargs) - end_time = time.time() + end_time = time.perf_counter_ns() # Calculate execution time execution_time = end_time - start_time # Measure overhead - overhead_start_time = time.time() - overhead_time = 0 + overhead_start_time = time.perf_counter_ns() try: # Check if currently in pytest benchmark fixture @@ -63,15 +62,16 @@ def wrapper(*args, **kwargs): benchmark_function_name = os.environ.get("CODEFLASH_BENCHMARK_FUNCTION_NAME", "") benchmark_file_name = os.environ.get("CODEFLASH_BENCHMARK_FILE_NAME", "") benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "") - - # Calculate overhead time - overhead_end_time = time.time() - overhead_time = overhead_end_time - overhead_start_time - + # Get class name class_name = "" qualname = func.__qualname__ if "." in qualname: class_name = qualname.split(".")[0] + # Calculate overhead time + overhead_end_time = time.perf_counter_ns() + overhead_time = overhead_end_time - overhead_start_time + + self.function_calls_data.append( (func.__name__, class_name, func.__module__, func.__code__.co_filename, benchmark_function_name, benchmark_file_name, benchmark_line_number, execution_time, @@ -84,114 +84,8 @@ def wrapper(*args, **kwargs): return result return wrapper - def write_to_db(self, output_file: str) -> None: - """Write all collected function call data to the SQLite database. - Args: - output_file: Path to the SQLite database file where results will be stored - """ - if not self.function_calls_data: - print("No function call data to write") - return - self.db_path = output_file - 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, module_name TEXT, file_name TEXT," - "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," - "time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" - ) - - # Insert all data at once - cur.executemany( - "INSERT INTO function_calls " - "(function_name, class_name, module_name, file_name, benchmark_function_name, " - "benchmark_file_name, benchmark_line_number, time_ns, overhead_time_ns, args, kwargs) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - self.function_calls_data - ) - - # Commit and close - con.commit() - con.close() - - print(f"Successfully wrote {len(self.function_calls_data)} function call records to {output_file}") - - # Clear the data after writing - self.function_calls_data.clear() - - except Exception as e: - print(f"Error writing function calls to database: {e}") - - def print_codeflash_db(self, limit: int = None) -> None: - """ - Print the contents of a CodeflashTrace SQLite database. - - Args: - db_path: Path to the SQLite database file - limit: Maximum number of records to print (None for all) - """ - try: - # Connect to the database - conn = sqlite3.connect(self.db_path) - cursor = conn.cursor() - - # Get the count of records - cursor.execute("SELECT COUNT(*) FROM function_calls") - total_records = cursor.fetchone()[0] - print(f"Found {total_records} function call records in {self.db_path}") - - # Build the query with optional limit - query = "SELECT * FROM function_calls" - if limit: - query += f" LIMIT {limit}" - - # Execute the query - cursor.execute(query) - - # Print column names - columns = [desc[0] for desc in cursor.description] - print("\nColumns:", columns) - print("\n" + "=" * 80 + "\n") - - # Print each row - for i, row in enumerate(cursor.fetchall()): - print(f"Record #{i + 1}:") - print(f" Function: {row[0]}") - print(f" Class: {row[1]}") - print(f" Module: {row[2]}") - print(f" File: {row[3]}") - print(f" Benchmark Function: {row[4] or 'N/A'}") - print(f" Benchmark File: {row[5] or 'N/A'}") - print(f" Benchmark Line: {row[6] or 'N/A'}") - print(f" Execution Time: {row[7]:.6f} seconds") - print(f" Overhead Time: {row[8]:.6f} seconds") - - # Unpickle and print args and kwargs - try: - args = pickle.loads(row[9]) - kwargs = pickle.loads(row[10]) - - print(f" Args: {args}") - print(f" Kwargs: {kwargs}") - except Exception as e: - print(f" Error unpickling args/kwargs: {e}") - print(f" Raw args: {row[8]}") - print(f" Raw kwargs: {row[9]}") - - print("\n" + "-" * 40 + "\n") - - conn.close() - - except Exception as e: - print(f"Error reading database: {e}") # Create a singleton instance diff --git a/codeflash/benchmarking/get_trace_info.py b/codeflash/benchmarking/get_trace_info.py index 3dd3831ce..e9a050b84 100644 --- a/codeflash/benchmarking/get_trace_info.py +++ b/codeflash/benchmarking/get_trace_info.py @@ -1,114 +1,109 @@ import sqlite3 from pathlib import Path -from typing import Dict, Set from codeflash.discovery.functions_to_optimize import FunctionToOptimize -def get_function_benchmark_timings(trace_dir: Path, all_functions_to_optimize: list[FunctionToOptimize]) -> dict[str, dict[str, float]]: - """Process all trace files in the given directory and extract timing data for the specified functions. +def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, float]]: + """Process the trace file and extract timing data for all functions. Args: - trace_dir: Path to the directory containing .trace files - all_functions_to_optimize: Set of FunctionToOptimize objects representing functions to include + trace_path: Path to the trace file + all_functions_to_optimize: List of FunctionToOptimize objects (not used directly, + but kept for backward compatibility) Returns: A nested dictionary where: - - Outer keys are function qualified names with file name - - Inner keys are benchmark names (trace filename without .trace extension) + - Outer keys are module_name.qualified_name (module.class.function) + - Inner keys are benchmark filename :: benchmark test function :: line number - Values are function timing in milliseconds """ - # Create a mapping of (filename, function_name, class_name) -> qualified_name for efficient lookups - function_lookup = {} - function_benchmark_timings = {} + # Initialize the result dictionary + result = {} + + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + + try: + # Query the function_calls table for all function calls + cursor.execute( + "SELECT module_name, class_name, function_name, " + "benchmark_file_name, benchmark_function_name, benchmark_line_number, " + "(time_ns - overhead_time_ns) as actual_time_ns " + "FROM function_calls" + ) + + # Process each row + for row in cursor.fetchall(): + module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row + + # Create the function key (module_name.class_name.function_name) + if class_name: + qualified_name = f"{module_name}.{class_name}.{function_name}" + else: + qualified_name = f"{module_name}.{function_name}" - for func in all_functions_to_optimize: - qualified_name = func.qualified_name_with_file_name + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - # Extract components (assumes Path.name gives only filename without directory) - filename = func.file_path - function_name = func.function_name + # Initialize the inner dictionary if needed + if qualified_name not in result: + result[qualified_name] = {} - # Get class name if there's a parent - class_name = func.parents[0].name if func.parents else None + # If multiple calls to the same function in the same benchmark, + # add the times together + if benchmark_key in result[qualified_name]: + result[qualified_name][benchmark_key] += time_ns + else: + result[qualified_name][benchmark_key] = time_ns - # Store in lookup dictionary - key = (filename, function_name, class_name) - function_lookup[key] = qualified_name - function_benchmark_timings[qualified_name] = {} + finally: + # Close the connection + connection.close() - # Find all .trace files in the directory - trace_files = list(trace_dir.glob("*.trace")) + return result - for trace_file in trace_files: - # Extract benchmark name from filename (without .trace) - benchmark_name = trace_file.stem - # Connect to the trace database - conn = sqlite3.connect(trace_file) - cursor = conn.cursor() +def get_benchmark_timings(trace_path: Path) -> dict[str, float]: + """Extract total benchmark timings from trace files. - # For each function we're interested in, query the database directly - for (filename, function_name, class_name), qualified_name in function_lookup.items(): - # Adjust query based on whether we have a class name - if class_name: - cursor.execute( - "SELECT cumulative_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND class_name = ?", - (f"%{filename}", function_name, class_name) - ) - else: - cursor.execute( - "SELECT cumulative_time_ns FROM pstats WHERE filename LIKE ? AND function = ? AND (class_name IS NULL OR class_name = '')", - (f"%{filename}", function_name) - ) + Args: + trace_path: Path to the trace file - result = cursor.fetchall() - if len(result) > 1: - print(f"Multiple results found for {qualified_name} in {benchmark_name}: {result}") - if result: - time_ns = result[0][0] - function_benchmark_timings[qualified_name][benchmark_name] = time_ns / 1e6 # Convert to milliseconds + Returns: + A dictionary mapping where: + - Keys are benchmark filename :: benchmark test function :: line number + - Values are total benchmark timing in milliseconds - conn.close() + """ + # Initialize the result dictionary + result = {} - return function_benchmark_timings + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + try: + # Query the benchmark_timings table + cursor.execute( + "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns " + "FROM benchmark_timings" + ) -def get_benchmark_timings(trace_dir: Path) -> dict[str, float]: - """Extract total benchmark timings from trace files. + # Process each row + for row in cursor.fetchall(): + benchmark_file, benchmark_func, benchmark_line, time_ns = row - Args: - trace_dir: Path to the directory containing .trace files + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - Returns: - A dictionary mapping benchmark names to their total execution time in milliseconds. - """ - benchmark_timings = {} - - # Find all .trace files in the directory - trace_files = list(trace_dir.glob("*.trace")) - - for trace_file in trace_files: - # Extract benchmark name from filename (without .trace extension) - benchmark_name = trace_file.stem - - # Connect to the trace database - conn = sqlite3.connect(trace_file) - cursor = conn.cursor() - - # Query the total_time table for the benchmark's total execution time - try: - cursor.execute("SELECT time_ns FROM total_time") - result = cursor.fetchone() - if result: - time_ns = result[0] - # Convert nanoseconds to milliseconds - benchmark_timings[benchmark_name] = time_ns / 1e6 - except sqlite3.OperationalError: - # Handle case where total_time table might not exist - print(f"Warning: Could not get total time for benchmark {benchmark_name}") - - conn.close() - - return benchmark_timings + # Store the timing + result[benchmark_key] = time_ns + + finally: + # Close the connection + connection.close() + + return result diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py index 99b2dad20..93f51baed 100644 --- a/codeflash/benchmarking/instrument_codeflash_trace.py +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -1,3 +1,4 @@ +import isort import libcst as cst from codeflash.discovery.functions_to_optimize import FunctionToOptimize @@ -9,6 +10,7 @@ def __init__(self, function_name, class_name=None): self.function_name = function_name self.class_name = class_name self.in_target_class = (class_name is None) # If no class name, always "in target class" + self.added_codeflash_trace = False def leave_ClassDef(self, original_node, updated_node): if self.class_name and original_node.name.value == self.class_name: @@ -31,12 +33,39 @@ def leave_FunctionDef(self, original_node, updated_node): # Add the new decorator after any existing decorators updated_decorators = list(updated_node.decorators) + [decorator] - + self.added_codeflash_trace = True # Return the updated node with the new decorator return updated_node.with_changes( decorators=updated_decorators ) + def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: + # Create import statement for codeflash_trace + if not self.added_codeflash_trace: + return updated_node + import_stmt = cst.SimpleStatementLine( + body=[ + cst.ImportFrom( + module=cst.Attribute( + value=cst.Attribute( + value=cst.Name(value="codeflash"), + attr=cst.Name(value="benchmarking") + ), + attr=cst.Name(value="codeflash_trace") + ), + names=[ + cst.ImportAlias( + name=cst.Name(value="codeflash_trace") + ) + ] + ) + ] + ) + + # Insert at the beginning of the file + new_body = [import_stmt, *list(updated_node.body)] + + return updated_node.with_changes(body=new_body) def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToOptimize) -> str: """Add codeflash_trace to a function. @@ -63,7 +92,7 @@ def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToO return modified_module.code -def instrument_codeflash_trace( +def instrument_codeflash_trace_decorator( function_to_optimize: FunctionToOptimize ) -> None: """Instrument __init__ function with codeflash_trace decorator if it's in a class.""" @@ -71,10 +100,10 @@ def instrument_codeflash_trace( original_code = function_to_optimize.file_path.read_text(encoding="utf-8") # Modify the code - modified_code = add_codeflash_decorator_to_code( + modified_code = isort.code(add_codeflash_decorator_to_code( original_code, function_to_optimize - ) + )) # Write the modified code back to the file function_to_optimize.file_path.write_text(modified_code, encoding="utf-8") diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index caf175a4e..a5f82fc3a 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -3,7 +3,8 @@ import pytest import time import os -class CodeFlashPlugin: +class CodeFlashBenchmarkPlugin: + benchmark_timings = [] @staticmethod def pytest_addoption(parser): parser.addoption( @@ -38,15 +39,20 @@ def benchmark(request): class Benchmark: def __call__(self, func, *args, **kwargs): - os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = request.node.name - os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = request.node.fspath.basename - os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack + benchmark_file_name = request.node.fspath.basename + benchmark_function_name = request.node.name + line_number = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack + os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name + os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = benchmark_file_name + os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = line_number os.environ["CODEFLASH_BENCHMARKING"] = "True" - start = time.process_time_ns() + + start = time.perf_counter_ns() result = func(*args, **kwargs) - end = time.process_time_ns() + end = time.perf_counter_ns() + os.environ["CODEFLASH_BENCHMARKING"] = "False" - print(f"Benchmark: {func.__name__} took {end - start} ns") + CodeFlashBenchmarkPlugin.benchmark_timings.append((benchmark_file_name, benchmark_function_name, line_number, end - start)) return result return Benchmark() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 04c5e67ea..a83196758 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -1,25 +1,31 @@ import sys from pathlib import Path +from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils from codeflash.verification.verification_utils import get_test_file_path -from plugin.plugin import CodeFlashPlugin +from plugin.plugin import CodeFlashBenchmarkPlugin from codeflash.benchmarking.codeflash_trace import codeflash_trace from codeflash.code_utils.code_utils import get_run_tmp_file benchmarks_root = sys.argv[1] tests_root = sys.argv[2] -output_file = sys.argv[3] +trace_file = sys.argv[3] # current working directory project_root = Path.cwd() if __name__ == "__main__": import pytest try: + db = BenchmarkDatabaseUtils(trace_path=Path(trace_file)) + db.setup() exitcode = pytest.main( - [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashPlugin()] + [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashBenchmarkPlugin()] ) - codeflash_trace.write_to_db(output_file) - codeflash_trace.print_codeflash_db() + db.write_function_timings(codeflash_trace.function_calls_data) + db.write_benchmark_timings(CodeFlashBenchmarkPlugin.benchmark_timings) + db.print_function_timings() + db.print_benchmark_timings() + db.close() except Exception as e: print(f"Failed to collect tests: {e!s}") diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 58ce456c2..75ef7e96d 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -270,8 +270,9 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework output_file = get_test_file_path( test_dir=Path(output_dir), function_name=f"{benchmark_file_name[5:]}_{benchmark_function_name}", test_type="replay" ) - with open(output_file, 'w') as f: - f.write(test_code) + # Write test code to file, parents = true + output_dir.mkdir(parents=True, exist_ok=True) + output_file.write_text(test_code, "utf-8") print(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") conn.close() diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 5c0a077dc..9ae69495d 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -3,19 +3,21 @@ from pathlib import Path import subprocess -def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, output_file: Path) -> None: +def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, trace_file: Path) -> None: + # set up .trace databases result = subprocess.run( [ SAFE_SYS_EXECUTABLE, Path(__file__).parent / "pytest_new_process_trace_benchmarks.py", benchmarks_root, tests_root, - output_file, + trace_file, ], cwd=project_root, check=False, capture_output=True, text=True, + env={"PYTHONPATH": str(project_root)}, ) print("stdout:", result.stdout) print("stderr:", result.stderr) diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index 774571de3..cd1e53f9b 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -143,11 +143,11 @@ def qualified_name(self) -> str: def qualified_name_with_modules_from_root(self, project_root_path: Path) -> str: return f"{module_name_from_file_path(self.file_path, project_root_path)}.{self.qualified_name}" - - @property - def qualified_name_with_file_name(self) -> str: - class_name = self.parents[0].name if self.parents else None - return f"{self.file_path}:{(class_name + ':' if class_name else '')}{self.function_name}" + # + # @property + # def qualified_name_with_file_name(self) -> str: + # class_name = self.parents[0].name if self.parents else None + # return f"{self.file_path}:{(class_name + ':' if class_name else '')}{self.function_name}" def get_functions_to_optimize( @@ -363,23 +363,28 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None: for decorator in body_node.decorator_list ): self.is_classmethod = True - return - else: - # search if the class has a staticmethod with the same name and on the same line number - for body_node in node.body: - if ( - isinstance(body_node, ast.FunctionDef) - and body_node.name == self.function_name - and body_node.lineno in {self.line_no, self.line_no + 1} - and any( + elif any( isinstance(decorator, ast.Name) and decorator.id == "staticmethod" for decorator in body_node.decorator_list - ) - ): - self.is_staticmethod = True - self.is_top_level = True - self.class_name = node.name + ): + self.is_staticmethod = True return + # else: + # # search if the class has a staticmethod with the same name and on the same line number + # for body_node in node.body: + # if ( + # isinstance(body_node, ast.FunctionDef) + # and body_node.name == self.function_name + # # and body_node.lineno in {self.line_no, self.line_no + 1} + # and any( + # isinstance(decorator, ast.Name) and decorator.id == "staticmethod" + # for decorator in body_node.decorator_list + # ) + # ): + # self.is_staticmethod = True + # self.is_top_level = True + # self.class_name = node.name + # return return diff --git a/codeflash/discovery/pytest_new_process_discover_benchmarks.py b/codeflash/discovery/pytest_new_process_discover_benchmarks.py deleted file mode 100644 index 83175218b..000000000 --- a/codeflash/discovery/pytest_new_process_discover_benchmarks.py +++ /dev/null @@ -1,54 +0,0 @@ -import sys -from typing import Any - -# This script should not have any relation to the codeflash package, be careful with imports -cwd = sys.argv[1] -tests_root = sys.argv[2] -pickle_path = sys.argv[3] -collected_tests = [] -pytest_rootdir = None -sys.path.insert(1, str(cwd)) - - -class PytestCollectionPlugin: - def pytest_collection_finish(self, session) -> None: - global pytest_rootdir - collected_tests.extend(session.items) - pytest_rootdir = session.config.rootdir - - -def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, str]]: - test_results = [] - for test in pytest_tests: - test_class = None - if test.cls: - test_class = test.parent.name - - # Determine if this is a benchmark test by checking for the benchmark fixture - is_benchmark = hasattr(test, 'fixturenames') and 'benchmark' in test.fixturenames - test_type = 'benchmark' if is_benchmark else 'regular' - - test_results.append({ - "test_file": str(test.path), - "test_class": test_class, - "test_function": test.name, - "test_type": test_type - }) - return test_results - - -if __name__ == "__main__": - import pytest - - try: - exitcode = pytest.main( - [tests_root, "-pno:logging", "--collect-only", "-m", "not skip"], plugins=[PytestCollectionPlugin()] - ) - except Exception as e: - print(f"Failed to collect tests: {e!s}") - exitcode = -1 - tests = parse_pytest_collection_results(collected_tests) - import pickle - - with open(pickle_path, "wb") as f: - pickle.dump((exitcode, tests, pytest_rootdir), f, protocol=pickle.HIGHEST_PROTOCOL) diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 249251e34..34556784f 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -280,16 +280,23 @@ def optimize_function(self) -> Result[BestOptimization, str]: speedup = explanation.speedup # if self.args.benchmark: original_replay_timing = original_code_baseline.benchmarking_test_results.total_replay_test_runtime() - fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_file_name] - for benchmark_name, og_benchmark_timing in fto_benchmark_timings.items(): - print(f"Calculating speedup for benchmark {benchmark_name}") - total_benchmark_timing = self.total_benchmark_timings[benchmark_name] + fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)] + for benchmark_key, og_benchmark_timing in fto_benchmark_timings.items(): + # benchmark key is benchmark filename :: benchmark test function :: line number + try: + benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") + except ValueError: + print(f"Benchmark key {benchmark_key} is not in the expected format.") + continue + print(f"Calculating speedup for benchmark {benchmark_key}") + total_benchmark_timing = self.total_benchmark_timings[benchmark_key] # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values + print(f"Original benchmark timing: {total_benchmark_timing}") replay_speedup = original_replay_timing / best_optimization.replay_runtime - 1 print(f"Replay speedup: {replay_speedup}") expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / (replay_speedup + 1) * og_benchmark_timing print(f"Expected new benchmark timing: {expected_new_benchmark_timing}") - print(f"Original benchmark timing: {total_benchmark_timing}") + benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 print(f"Benchmark speedup: {benchmark_speedup_percent:.2f}%") diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 5f5f0ec2f..0b96e6b77 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient +from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.utils import print_benchmark_table from codeflash.cli_cmds.console import console, logger @@ -25,7 +26,8 @@ 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 codeflash.benchmarking.codeflash_trace import codeflash_trace +from codeflash.benchmarking.instrument_codeflash_trace import instrument_codeflash_trace_decorator + from collections import defaultdict @@ -94,6 +96,8 @@ def run(self) -> None: project_root=self.args.project_root, module_root=self.args.module_root, ) + all_functions_to_optimize = [ + fto for functions_to_optimize in file_to_funcs_to_optimize.values() for fto in functions_to_optimize] if self.args.benchmark: # Insert decorator file_path_to_source_code = defaultdict(str) @@ -103,9 +107,14 @@ def run(self) -> None: 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 + instrument_codeflash_trace_decorator(fto) + trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" + trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Simply run all tests that use pytest-benchmark + generate_replay_test(trace_file, Path(self.args.tests_root) / "codeflash_replay_tests" ) + function_benchmark_timings = get_function_benchmark_timings(trace_file) + total_benchmark_timings = get_benchmark_timings(trace_file) + print(function_benchmark_timings) + print(total_benchmark_timings) logger.info("Finished tracing existing benchmarks") except Exception as e: logger.info(f"Error while tracing existing benchmarks: {e}") @@ -116,13 +125,13 @@ def run(self) -> None: with file.open("w", encoding="utf8") as f: f.write(file_path_to_source_code[file]) - codeflash_trace.print_trace_info() + # 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) - + # return optimizations_found: int = 0 function_iterator_count: int = 0 if self.args.test_framework == "pytest": @@ -206,6 +215,10 @@ def run(self) -> None: function_optimizer = self.create_function_optimizer( function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code, function_benchmark_timings, total_benchmark_timings ) + # function_optimizer = self.create_function_optimizer( + # function_to_optimize, function_to_optimize_ast, function_to_tests, + # validated_original_code[original_module_path].source_code + # ) else: function_optimizer = self.create_function_optimizer( function_to_optimize, function_to_optimize_ast, function_to_tests, diff --git a/tests/test_instrument_codeflash_trace.py b/tests/test_instrument_codeflash_trace.py index 56008faa9..967d5d6f0 100644 --- a/tests/test_instrument_codeflash_trace.py +++ b/tests/test_instrument_codeflash_trace.py @@ -26,6 +26,7 @@ def normal_function(): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace @codeflash_trace def normal_function(): return "Hello, World!" @@ -53,6 +54,7 @@ def normal_method(self): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @codeflash_trace def normal_method(self): @@ -82,6 +84,7 @@ def class_method(cls): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @classmethod @codeflash_trace @@ -112,6 +115,7 @@ def static_method(): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @staticmethod @codeflash_trace @@ -141,6 +145,7 @@ def __init__(self, value): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @codeflash_trace def __init__(self, value): @@ -171,6 +176,7 @@ def property_method(self): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @property @other_decorator @@ -205,6 +211,7 @@ def test_method(self): ) expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace class TestClass: @codeflash_trace def test_method(self): From 54fe71f336e40b8ab472aaae7f4d9db3d4944cfd Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 19 Mar 2025 15:04:47 -0700 Subject: [PATCH 065/122] modified printing of results, handle errors when collecting benchmarks --- codeflash/benchmarking/codeflash_trace.py | 4 - codeflash/benchmarking/get_trace_info.py | 34 +++++--- .../pytest_new_process_trace_benchmarks.py | 11 +-- codeflash/benchmarking/replay_test.py | 13 +-- codeflash/benchmarking/trace_benchmarks.py | 25 +++++- codeflash/benchmarking/utils.py | 24 ++++-- codeflash/discovery/functions_to_optimize.py | 9 +- codeflash/models/models.py | 1 + codeflash/optimization/function_optimizer.py | 45 +++------- codeflash/optimization/optimizer.py | 83 +++++++++---------- codeflash/result/explanation.py | 41 +++++++-- codeflash/verification/test_runner.py | 2 + 12 files changed, 164 insertions(+), 128 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 9b9afead7..14505efee 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -84,9 +84,5 @@ def wrapper(*args, **kwargs): return result return wrapper - - - - # Create a singleton instance codeflash_trace = CodeflashTrace() diff --git a/codeflash/benchmarking/get_trace_info.py b/codeflash/benchmarking/get_trace_info.py index e9a050b84..d43327af7 100644 --- a/codeflash/benchmarking/get_trace_info.py +++ b/codeflash/benchmarking/get_trace_info.py @@ -4,13 +4,11 @@ from codeflash.discovery.functions_to_optimize import FunctionToOptimize -def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, float]]: +def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, int]]: """Process the trace file and extract timing data for all functions. Args: trace_path: Path to the trace file - all_functions_to_optimize: List of FunctionToOptimize objects (not used directly, - but kept for backward compatibility) Returns: A nested dictionary where: @@ -30,8 +28,7 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, floa # Query the function_calls table for all function calls cursor.execute( "SELECT module_name, class_name, function_name, " - "benchmark_file_name, benchmark_function_name, benchmark_line_number, " - "(time_ns - overhead_time_ns) as actual_time_ns " + "benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns " "FROM function_calls" ) @@ -66,7 +63,7 @@ def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, floa return result -def get_benchmark_timings(trace_path: Path) -> dict[str, float]: +def get_benchmark_timings(trace_path: Path) -> dict[str, int]: """Extract total benchmark timings from trace files. Args: @@ -75,32 +72,47 @@ def get_benchmark_timings(trace_path: Path) -> dict[str, float]: Returns: A dictionary mapping where: - Keys are benchmark filename :: benchmark test function :: line number - - Values are total benchmark timing in milliseconds + - Values are total benchmark timing in milliseconds (with overhead subtracted) """ # Initialize the result dictionary result = {} + overhead_by_benchmark = {} # Connect to the SQLite database connection = sqlite3.connect(trace_path) cursor = connection.cursor() try: - # Query the benchmark_timings table + # Query the function_calls table to get total overhead for each benchmark + cursor.execute( + "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " + "FROM function_calls " + "GROUP BY benchmark_file_name, 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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + 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_name, benchmark_function_name, benchmark_line_number, time_ns " "FROM benchmark_timings" ) - # Process each row + # Process each row and subtract overhead for row in cursor.fetchall(): benchmark_file, benchmark_func, benchmark_line, time_ns = row # Create the benchmark key (file::function::line) benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - # Store the timing - result[benchmark_key] = time_ns + # Subtract overhead from total time + overhead = overhead_by_benchmark.get(benchmark_key, 0) + result[benchmark_key] = time_ns - overhead finally: # Close the connection diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index a83196758..6d4c85f41 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -20,13 +20,14 @@ db.setup() exitcode = pytest.main( [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashBenchmarkPlugin()] - ) + ) # Errors will be printed to stdout, not stderr db.write_function_timings(codeflash_trace.function_calls_data) db.write_benchmark_timings(CodeFlashBenchmarkPlugin.benchmark_timings) - db.print_function_timings() - db.print_benchmark_timings() + # db.print_function_timings() + # db.print_benchmark_timings() db.close() except Exception as e: - print(f"Failed to collect tests: {e!s}") - exitcode = -1 \ No newline at end of file + print(f"Failed to collect tests: {e!s}", file=sys.stderr) + exitcode = -1 + sys.exit(exitcode) \ No newline at end of file diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 75ef7e96d..a1d5b370a 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -142,8 +142,6 @@ def create_trace_replay_test_code( class_name = func.get("class_name") file_name = func.get("file_name") function_properties = func.get("function_properties") - print(f"Class: {class_name}, Function: {function_name}") - print(function_properties) if not class_name: alias = get_function_alias(module_name, function_name) test_body = test_function_body.format( @@ -197,7 +195,7 @@ def create_trace_replay_test_code( return imports + "\n" + metadata + "\n" + test_template -def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework: str = "pytest", max_run_count: int = 100) -> None: +def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework: str = "pytest", max_run_count: int = 100) -> int: """Generate multiple replay tests from the traced function calls, grouping by benchmark name. Args: @@ -211,6 +209,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework Dictionary mapping benchmark names to generated test code """ + count = 0 try: # Connect to the database conn = sqlite3.connect(trace_file_path.as_posix()) @@ -253,7 +252,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework }) if not functions_data: - print(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") + logger.info(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") continue # Generate the test code for this benchmark @@ -273,9 +272,11 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework # Write test code to file, parents = true output_dir.mkdir(parents=True, exist_ok=True) output_file.write_text(test_code, "utf-8") - print(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") + count += 1 + logger.info(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") conn.close() except Exception as e: - print(f"Error generating replay tests: {e}") + logger.info(f"Error generating replay tests: {e}") + return count \ No newline at end of file diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 9ae69495d..79395db79 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -1,10 +1,15 @@ from __future__ import annotations + +import re + +from pytest import ExitCode + +from codeflash.cli_cmds.console import logger from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE from pathlib import Path import subprocess def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, trace_file: Path) -> None: - # set up .trace databases result = subprocess.run( [ SAFE_SYS_EXECUTABLE, @@ -19,5 +24,19 @@ def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root text=True, env={"PYTHONPATH": str(project_root)}, ) - print("stdout:", result.stdout) - print("stderr:", result.stderr) + if result.returncode != 0: + if "ERROR collecting" in result.stdout: + # Pattern matches "===== ERRORS =====" (any number of =) and captures everything after + error_pattern = r"={3,}\s*ERRORS\s*={3,}\n([\s\S]*?)(?:={3,}|$)" + match = re.search(error_pattern, result.stdout) + error_section = match.group(1) if match else result.stdout + elif "FAILURES" in result.stdout: + # Pattern matches "===== FAILURES =====" (any number of =) and captures everything after + error_pattern = r"={3,}\s*FAILURES\s*={3,}\n([\s\S]*?)(?:={3,}|$)" + match = re.search(error_pattern, result.stdout) + error_section = match.group(1) if match else result.stdout + else: + error_section = result.stdout + logger.warning( + f"Error collecting benchmarks - Pytest Exit code: {result.returncode}={ExitCode(result.returncode).name}\n {error_section}" + ) \ No newline at end of file diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index d97c2e36e..685bfe739 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,7 +1,12 @@ -def print_benchmark_table(function_benchmark_timings, total_benchmark_timings): +def print_benchmark_table(function_benchmark_timings: dict[str,dict[str,int]], total_benchmark_timings: dict[str,int]): + # Define column widths + benchmark_col_width = 50 + time_col_width = 15 + # Print table header - print(f"{'Benchmark Test':<50} | {'Total Time (s)':<15} | {'Function Time (s)':<15} | {'Percentage (%)':<15}") - print("-" * 100) + header = f"{'Benchmark Test':{benchmark_col_width}} | {'Total Time (ms)':{time_col_width}} | {'Function Time (ms)':{time_col_width}} | {'Percentage (%)':{time_col_width}}" + print(header) + print("-" * len(header)) # Process each function's benchmark data for func_path, test_times in function_benchmark_timings.items(): @@ -14,13 +19,16 @@ def print_benchmark_table(function_benchmark_timings, total_benchmark_timings): total_time = total_benchmark_timings.get(test_name, 0) if total_time > 0: percentage = (func_time / total_time) * 100 - sorted_tests.append((test_name, total_time, func_time, percentage)) + # Convert nanoseconds to milliseconds + func_time_ms = func_time / 1_000_000 + total_time_ms = total_time / 1_000_000 + sorted_tests.append((test_name, total_time_ms, func_time_ms, percentage)) sorted_tests.sort(key=lambda x: x[3], reverse=True) # Print each test's data for test_name, total_time, func_time, percentage in sorted_tests: - print(f"{test_name:<50} | {total_time:<15.3f} | {func_time:<15.3f} | {percentage:<15.2f}") - -# Usage - + benchmark_file, benchmark_func, benchmark_line = test_name.split("::") + benchmark_name = f"{benchmark_file}::{benchmark_func}" + print(f"{benchmark_name:{benchmark_col_width}} | {total_time:{time_col_width}.3f} | {func_time:{time_col_width}.3f} | {percentage:{time_col_width}.2f}") + print() diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index cd1e53f9b..fb80541aa 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -106,7 +106,7 @@ def generic_visit(self, node: ast.AST) -> None: @dataclass(frozen=True, config={"arbitrary_types_allowed": True}) class FunctionToOptimize: - """Represents a function that is a candidate for optimization. + """Represent a function that is a candidate for optimization. Attributes ---------- @@ -121,6 +121,7 @@ class FunctionToOptimize: method extends this with the module name from the project root. """ + function_name: str file_path: Path parents: list[FunctionParent] # list[ClassDef | FunctionDef | AsyncFunctionDef] @@ -143,12 +144,6 @@ def qualified_name(self) -> str: def qualified_name_with_modules_from_root(self, project_root_path: Path) -> str: return f"{module_name_from_file_path(self.file_path, project_root_path)}.{self.qualified_name}" - # - # @property - # def qualified_name_with_file_name(self) -> str: - # class_name = self.parents[0].name if self.parents else None - # return f"{self.file_path}:{(class_name + ':' if class_name else '')}{self.function_name}" - def get_functions_to_optimize( optimize_all: str | None, diff --git a/codeflash/models/models.py b/codeflash/models/models.py index c0ce74c47..5358b8e4e 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -77,6 +77,7 @@ class BestOptimization(BaseModel): helper_functions: list[FunctionSource] runtime: int replay_runtime: int | None + replay_performance_gain: float | None winning_behavioral_test_results: TestResults winning_benchmarking_test_results: TestResults winning_replay_benchmarking_test_results : TestResults | None = None diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 34556784f..b2af5e0bf 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -92,8 +92,8 @@ def __init__( function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_ast: ast.FunctionDef | None = None, aiservice_client: AiServiceClient | None = None, - function_benchmark_timings: dict[str, dict[str, float]] | None = None, - total_benchmark_timings: dict[str, float] | None = None, + function_benchmark_timings: dict[str, dict[str, int]] | None = None, + total_benchmark_timings: dict[str, int] | None = None, args: Namespace | None = None, ) -> None: self.project_root = test_cfg.project_root_path @@ -276,30 +276,10 @@ def optimize_function(self) -> Result[BestOptimization, str]: best_runtime_ns=best_optimization.runtime, function_name=function_to_optimize_qualified_name, file_path=self.function_to_optimize.file_path, + replay_performance_gain=best_optimization.replay_performance_gain if self.args.benchmark else None, + fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)] if self.args.benchmark else None, + total_benchmark_timings = self.total_benchmark_timings if self.args.benchmark else None, ) - speedup = explanation.speedup # - if self.args.benchmark: - original_replay_timing = original_code_baseline.benchmarking_test_results.total_replay_test_runtime() - fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)] - for benchmark_key, og_benchmark_timing in fto_benchmark_timings.items(): - # benchmark key is benchmark filename :: benchmark test function :: line number - try: - benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") - except ValueError: - print(f"Benchmark key {benchmark_key} is not in the expected format.") - continue - print(f"Calculating speedup for benchmark {benchmark_key}") - total_benchmark_timing = self.total_benchmark_timings[benchmark_key] - # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values - print(f"Original benchmark timing: {total_benchmark_timing}") - replay_speedup = original_replay_timing / best_optimization.replay_runtime - 1 - print(f"Replay speedup: {replay_speedup}") - expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / (replay_speedup + 1) * og_benchmark_timing - print(f"Expected new benchmark timing: {expected_new_benchmark_timing}") - - benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing - benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 - print(f"Benchmark speedup: {benchmark_speedup_percent:.2f}%") self.log_successful_optimization(explanation, generated_tests) @@ -455,21 +435,21 @@ def determine_best_candidate( original_runtime_ns=original_code_replay_runtime, optimized_runtime_ns=candidate_replay_runtime, ) - tree.add("Replay Benchmarking： ") - tree.add(f"Original summed runtime: {humanize_runtime(original_code_replay_runtime)}") + tree.add(f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") tree.add( - f"Best summed runtime: {humanize_runtime(candidate_replay_runtime)} " + f"Best benchmark replay runtime: {humanize_runtime(candidate_replay_runtime)} " f"(measured over {candidate_result.max_loop_count} " f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" ) - tree.add(f"Speedup percentage: {replay_perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio: {replay_perf_gain + 1:.1f}X") + tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain + 1:.1f}X") best_optimization = BestOptimization( candidate=candidate, helper_functions=code_context.helper_functions, runtime=best_test_runtime, replay_runtime=candidate_replay_runtime if self.args.benchmark else None, winning_behavioral_test_results=candidate_result.behavior_test_results, + replay_performance_gain=replay_perf_gain if self.args.benchmark else None, winning_benchmarking_test_results=candidate_result.benchmarking_test_results, winning_replay_benchmarking_test_results=candidate_result.benchmarking_test_results, ) @@ -525,7 +505,8 @@ def log_successful_optimization(self, explanation: Explanation, generated_tests: ) console.print(Group(explanation_panel, tests_panel)) - console.print(explanation_panel) + else: + console.print(explanation_panel) ph( "cli-optimize-success", @@ -682,7 +663,6 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi existing_test_files_count += 1 elif test_type == TestType.REPLAY_TEST: replay_test_files_count += 1 - print("Replay test found") elif test_type == TestType.CONCOLIC_COVERAGE_TEST: concolic_coverage_test_files_count += 1 else: @@ -1157,7 +1137,6 @@ def run_and_parse_tests( f"stdout: {run_result.stdout}\n" f"stderr: {run_result.stderr}\n" ) - # print(test_files) results, coverage_results = parse_test_results( test_xml_path=result_file_path, test_files=test_files, diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 0b96e6b77..894a8911e 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -11,8 +11,9 @@ from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.utils import print_benchmark_table -from codeflash.cli_cmds.console import console, logger +from codeflash.cli_cmds.console import console, logger, progress_bar from codeflash.code_utils import env_utils +from codeflash.code_utils.code_extractor import add_needed_imports_from_module from codeflash.code_utils.code_replacer import normalize_code, normalize_node from codeflash.code_utils.code_utils import get_run_tmp_file from codeflash.code_utils.static_analysis import analyze_imported_modules, get_first_top_level_function_or_method_ast @@ -96,42 +97,43 @@ def run(self) -> None: project_root=self.args.project_root, module_root=self.args.module_root, ) - all_functions_to_optimize = [ - fto for functions_to_optimize in file_to_funcs_to_optimize.values() for fto in functions_to_optimize] + function_benchmark_timings = None + total_benchmark_timings = None if self.args.benchmark: - # 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: - instrument_codeflash_trace_decorator(fto) - trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" - trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Simply run all tests that use pytest-benchmark - generate_replay_test(trace_file, Path(self.args.tests_root) / "codeflash_replay_tests" ) - function_benchmark_timings = get_function_benchmark_timings(trace_file) - total_benchmark_timings = get_benchmark_timings(trace_file) - print(function_benchmark_timings) - print(total_benchmark_timings) - logger.info("Finished tracing existing benchmarks") - except Exception as e: - logger.info(f"Error while tracing existing benchmarks: {e}") - logger.info(f"Information on existing benchmarks will not be available for this run.") - 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]) - + with progress_bar( + f"Running benchmarks in {self.args.benchmarks_root}", + transient=True, + ): + # 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: + instrument_codeflash_trace_decorator(fto) + trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" + replay_tests_dir = Path(self.args.tests_root) / "codeflash_replay_tests" + trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Simply run all tests that use pytest-benchmark + replay_count = generate_replay_test(trace_file, replay_tests_dir) + if replay_count == 0: + logger.info(f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization") + else: + function_benchmark_timings = get_function_benchmark_timings(trace_file) + total_benchmark_timings = get_benchmark_timings(trace_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) + print_benchmark_table(function_benchmark_timings, total_benchmark_timings) + logger.info("Finished tracing existing benchmarks") + except Exception as e: + logger.info(f"Error while tracing existing benchmarks: {e}") + logger.info(f"Information on existing benchmarks will not be available for this run.") + 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]) - # return optimizations_found: int = 0 function_iterator_count: int = 0 if self.args.test_framework == "pytest": @@ -210,15 +212,10 @@ def run(self) -> None: f"Skipping optimization." ) continue - if self.args.benchmark: - + if self.args.benchmark and function_benchmark_timings and total_benchmark_timings: function_optimizer = self.create_function_optimizer( function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code, function_benchmark_timings, total_benchmark_timings ) - # function_optimizer = self.create_function_optimizer( - # function_to_optimize, function_to_optimize_ast, function_to_tests, - # validated_original_code[original_module_path].source_code - # ) else: function_optimizer = self.create_function_optimizer( function_to_optimize, function_to_optimize_ast, function_to_tests, @@ -250,9 +247,9 @@ def run(self) -> None: if function_optimizer.test_cfg.concolic_test_root_dir: shutil.rmtree(function_optimizer.test_cfg.concolic_test_root_dir, ignore_errors=True) if self.args.benchmark: - trace_dir = Path(self.args.benchmarks_root) / "codeflash_replay_tests" - if trace_dir.exists(): - shutil.rmtree(trace_dir, ignore_errors=True) + if replay_tests_dir.exists(): + shutil.rmtree(replay_tests_dir, ignore_errors=True) + trace_file.unlink(missing_ok=True) if hasattr(get_run_tmp_file, "tmpdir"): get_run_tmp_file.tmpdir.cleanup() diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 8a2f8f81d..43eb973f7 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -15,6 +15,9 @@ class Explanation: best_runtime_ns: int function_name: str file_path: Path + replay_performance_gain: float | None + fto_benchmark_timings: dict[str, int] | None + total_benchmark_timings: dict[str, int] | None @property def perf_improvement_line(self) -> str: @@ -37,16 +40,38 @@ def to_console_string(self) -> str: # TODO: Sometimes the explanation says something similar to "This is the code that was optimized", remove such parts original_runtime_human = humanize_runtime(self.original_runtime_ns) best_runtime_human = humanize_runtime(self.best_runtime_ns) + benchmark_info = "" + if self.replay_performance_gain: + benchmark_info += "Benchmark Performance Details:\n" + for benchmark_key, og_benchmark_timing in self.fto_benchmark_timings.items(): + # benchmark key is benchmark filename :: benchmark test function :: line number + try: + benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") + except ValueError: + benchmark_info += f"Benchmark key {benchmark_key} is not in the expected format.\n" + continue + + total_benchmark_timing = self.total_benchmark_timings[benchmark_key] + # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values + benchmark_info += f"Original timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(total_benchmark_timing)}\n" + replay_speedup = self.replay_performance_gain + expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / ( + replay_speedup + 1) * og_benchmark_timing + benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(expected_new_benchmark_timing)}\n" + + benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing + benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 + benchmark_info += f"Benchmark speedup for {benchmark_file_name}::{benchmark_test_function}: {benchmark_speedup_percent:.2f}%\n\n" return ( - f"Optimized {self.function_name} in {self.file_path}\n" - f"{self.perf_improvement_line}\n" - f"Runtime went down from {original_runtime_human} to {best_runtime_human} \n\n" - + "Explanation:\n" - + self.raw_explanation_message - + " \n\n" - + "The new optimized code was tested for correctness. The results are listed below.\n" - + f"{TestResults.report_to_string(self.winning_behavioral_test_results.get_test_pass_fail_report_by_type())}\n" + f"Optimized {self.function_name} in {self.file_path}\n" + f"{self.perf_improvement_line}\n" + f"Runtime went down from {original_runtime_human} to {best_runtime_human} \n\n" + + (benchmark_info if benchmark_info else "") + + self.raw_explanation_message + + " \n\n" + + "The new optimized code was tested for correctness. The results are listed below.\n" + + f"{TestResults.report_to_string(self.winning_behavioral_test_results.get_test_pass_fail_report_by_type())}\n" ) def explanation_message(self) -> str: diff --git a/codeflash/verification/test_runner.py b/codeflash/verification/test_runner.py index 852b6bf8a..ab17c94a0 100644 --- a/codeflash/verification/test_runner.py +++ b/codeflash/verification/test_runner.py @@ -59,6 +59,7 @@ def run_behavioral_tests( ) test_files = list(set(test_files)) # remove multiple calls in the same test function common_pytest_args = [ + "--benchmark-skip", "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", @@ -240,6 +241,7 @@ def run_benchmarking_tests( test_files.append(str(file.benchmarking_file_path)) test_files = list(set(test_files)) # remove multiple calls in the same test function pytest_args = [ + "--benchmark-skip", "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", From 5fd112a2f9e633aa12f0cf2c997051b147da9f31 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 19 Mar 2025 15:59:13 -0700 Subject: [PATCH 066/122] tests pass --- code_to_optimize/bubble_sort.py | 2 ++ ...process_and_bubble_sort_codeflash_trace.py | 28 +++++++++++++++ codeflash/discovery/functions_to_optimize.py | 34 ++++++++++--------- .../discovery/pytest_new_process_discovery.py | 8 +---- tests/test_unit_test_discovery.py | 9 ++--- 5 files changed, 52 insertions(+), 29 deletions(-) create mode 100644 code_to_optimize/process_and_bubble_sort_codeflash_trace.py diff --git a/code_to_optimize/bubble_sort.py b/code_to_optimize/bubble_sort.py index db7db5f92..9e97f63a0 100644 --- a/code_to_optimize/bubble_sort.py +++ b/code_to_optimize/bubble_sort.py @@ -1,8 +1,10 @@ def sorter(arr): + print("codeflash stdout: Sorting list") for i in range(len(arr)): for j in range(len(arr) - 1): if arr[j] > arr[j + 1]: temp = arr[j] arr[j] = arr[j + 1] arr[j + 1] = temp + print(f"result: {arr}") return arr diff --git a/code_to_optimize/process_and_bubble_sort_codeflash_trace.py b/code_to_optimize/process_and_bubble_sort_codeflash_trace.py new file mode 100644 index 000000000..37c2abab8 --- /dev/null +++ b/code_to_optimize/process_and_bubble_sort_codeflash_trace.py @@ -0,0 +1,28 @@ +from code_to_optimize.bubble_sort import sorter +from codeflash.benchmarking.codeflash_trace import codeflash_trace + +def calculate_pairwise_products(arr): + """ + Calculate the average of all pairwise products in the array. + """ + sum_of_products = 0 + count = 0 + + for i in range(len(arr)): + for j in range(len(arr)): + if i != j: + sum_of_products += arr[i] * arr[j] + count += 1 + + # The average of all pairwise products + return sum_of_products / count if count > 0 else 0 + +@codeflash_trace +def compute_and_sort(arr): + # Compute pairwise sums average + pairwise_average = calculate_pairwise_products(arr) + + # Call sorter function + sorter(arr.copy()) + + return pairwise_average diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index fb80541aa..cd0bfc50a 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -363,23 +363,25 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None: for decorator in body_node.decorator_list ): self.is_staticmethod = True + print(f"static method found: {self.function_name}") + return + elif self.line_no: + # If we have line number info, check if class has a static method with the same line number + # This way, if we don't have the class name, we can still find the static method + for body_node in node.body: + if ( + isinstance(body_node, ast.FunctionDef) + and body_node.name == self.function_name + and body_node.lineno in {self.line_no, self.line_no + 1} + and any( + isinstance(decorator, ast.Name) and decorator.id == "staticmethod" + for decorator in body_node.decorator_list + ) + ): + self.is_staticmethod = True + self.is_top_level = True + self.class_name = node.name return - # else: - # # search if the class has a staticmethod with the same name and on the same line number - # for body_node in node.body: - # if ( - # isinstance(body_node, ast.FunctionDef) - # and body_node.name == self.function_name - # # and body_node.lineno in {self.line_no, self.line_no + 1} - # and any( - # isinstance(decorator, ast.Name) and decorator.id == "staticmethod" - # for decorator in body_node.decorator_list - # ) - # ): - # self.is_staticmethod = True - # self.is_top_level = True - # self.class_name = node.name - # return return diff --git a/codeflash/discovery/pytest_new_process_discovery.py b/codeflash/discovery/pytest_new_process_discovery.py index 2d8583255..d5a80f501 100644 --- a/codeflash/discovery/pytest_new_process_discovery.py +++ b/codeflash/discovery/pytest_new_process_discovery.py @@ -16,12 +16,6 @@ def pytest_collection_finish(self, session) -> None: collected_tests.extend(session.items) pytest_rootdir = session.config.rootdir - def pytest_collection_modifyitems(config, items): - skip_benchmark = pytest.mark.skip(reason="Skipping benchmark tests") - for item in items: - if "benchmark" in item.fixturenames: - item.add_marker(skip_benchmark) - def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, str]]: test_results = [] @@ -40,7 +34,7 @@ def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, s try: exitcode = pytest.main( - [tests_root, "-p no:logging", "--collect-only", "-m", "not skip"], plugins=[PytestCollectionPlugin()] + [tests_root, "-pno:logging", "--collect-only", "-m", "not skip", "--benchmark-skip"], plugins=[PytestCollectionPlugin()] ) except Exception as e: # noqa: BLE001 print(f"Failed to collect tests: {e!s}") # noqa: T201 diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index 4bf99c049..c05b79e63 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -18,11 +18,11 @@ def test_unit_test_discovery_pytest(): ) tests = discover_unit_tests(test_config) assert len(tests) > 0 - # print(tests) + def test_benchmark_test_discovery_pytest(): project_path = Path(__file__).parent.parent.resolve() / "code_to_optimize" - tests_path = project_path / "tests" / "pytest" / "benchmarks" / "test_benchmark_bubble_sort.py" + tests_path = project_path / "tests" / "pytest" / "benchmarks" test_config = TestConfig( tests_root=tests_path, project_root_path=project_path, @@ -30,10 +30,7 @@ def test_benchmark_test_discovery_pytest(): tests_project_rootdir=tests_path.parent, ) tests = discover_unit_tests(test_config) - assert len(tests) > 0 - assert 'bubble_sort.sorter' in tests - benchmark_tests = sum(1 for test in tests['bubble_sort.sorter'] if test.tests_in_file.test_type == TestType.BENCHMARK_TEST) - assert benchmark_tests == 1 + assert len(tests) == 1 # Should not discover benchmark tests def test_unit_test_discovery_unittest(): From 8194554e1b4285ff933e2590b75321d962f8348f Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 19 Mar 2025 16:00:45 -0700 Subject: [PATCH 067/122] revert pyproject.toml --- pyproject.toml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 877815004..2e71f2a0a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -216,9 +216,8 @@ initial-content = """ [tool.codeflash] -module-root = "code_to_optimize" -tests-root = "code_to_optimize/tests" -benchmarks-root = "code_to_optimize/tests/pytest/benchmarks" +module-root = "codeflash" +tests-root = "tests" test-framework = "pytest" formatter-cmds = [ "uvx ruff check --exit-zero --fix $file", From 4c1d2aff391e8f94dea12a8eb252e1070339306e Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 19 Mar 2025 16:02:44 -0700 Subject: [PATCH 068/122] mypy fixes --- codeflash/result/explanation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 43eb973f7..77528a8a6 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -41,7 +41,7 @@ def to_console_string(self) -> str: original_runtime_human = humanize_runtime(self.original_runtime_ns) best_runtime_human = humanize_runtime(self.best_runtime_ns) benchmark_info = "" - if self.replay_performance_gain: + if self.replay_performance_gain and self.fto_benchmark_timings and self.total_benchmark_timings: benchmark_info += "Benchmark Performance Details:\n" for benchmark_key, og_benchmark_timing in self.fto_benchmark_timings.items(): # benchmark key is benchmark filename :: benchmark test function :: line number @@ -57,7 +57,7 @@ def to_console_string(self) -> str: replay_speedup = self.replay_performance_gain expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / ( replay_speedup + 1) * og_benchmark_timing - benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(expected_new_benchmark_timing)}\n" + benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(int(expected_new_benchmark_timing))}\n" benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 From 6e676e90b2d45bbe5004d4e6c0b535d45ce92749 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 10:10:59 -0700 Subject: [PATCH 069/122] import changes --- codeflash/benchmarking/codeflash_trace.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 14505efee..b1236ffbf 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -1,9 +1,7 @@ import functools import os import pickle -import sqlite3 import time -from pathlib import Path from typing import Callable From 62f3b368e0ab121a6bdc0eefb64846faed55fa72 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 13:45:36 -0700 Subject: [PATCH 070/122] removed benchmark skip command --- codeflash/verification/test_runner.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/codeflash/verification/test_runner.py b/codeflash/verification/test_runner.py index ab17c94a0..852b6bf8a 100644 --- a/codeflash/verification/test_runner.py +++ b/codeflash/verification/test_runner.py @@ -59,7 +59,6 @@ def run_behavioral_tests( ) test_files = list(set(test_files)) # remove multiple calls in the same test function common_pytest_args = [ - "--benchmark-skip", "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", @@ -241,7 +240,6 @@ def run_benchmarking_tests( test_files.append(str(file.benchmarking_file_path)) test_files = list(set(test_files)) # remove multiple calls in the same test function pytest_args = [ - "--benchmark-skip", "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", From a6149726640c3e8a3cf7d6a9e79b94a7ba3d8bec Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 15:09:21 -0700 Subject: [PATCH 071/122] shifted benchmark class in plugin, improved display of benchmark info --- codeflash/benchmarking/plugin/plugin.py | 44 ++++++++++++++----------- codeflash/benchmarking/utils.py | 37 ++++++++++++++------- 2 files changed, 49 insertions(+), 32 deletions(-) diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index a5f82fc3a..ee7504ec4 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -5,6 +5,29 @@ import os class CodeFlashBenchmarkPlugin: benchmark_timings = [] + + class Benchmark: + def __init__(self, request): + self.request = request + + def __call__(self, func, *args, **kwargs): + benchmark_file_name = self.request.node.fspath.basename + benchmark_function_name = self.request.node.name + line_number = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack + + os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name + os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = benchmark_file_name + os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = line_number + os.environ["CODEFLASH_BENCHMARKING"] = "True" + + start = time.perf_counter_ns() + result = func(*args, **kwargs) + end = time.perf_counter_ns() + + os.environ["CODEFLASH_BENCHMARKING"] = "False" + CodeFlashBenchmarkPlugin.benchmark_timings.append( + (benchmark_file_name, benchmark_function_name, line_number, end - start)) + return result @staticmethod def pytest_addoption(parser): parser.addoption( @@ -36,23 +59,4 @@ def benchmark(request): if not request.config.getoption("--codeflash-trace"): return None - class Benchmark: - - def __call__(self, func, *args, **kwargs): - benchmark_file_name = request.node.fspath.basename - benchmark_function_name = request.node.name - line_number = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack - os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name - os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = benchmark_file_name - os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = line_number - os.environ["CODEFLASH_BENCHMARKING"] = "True" - - start = time.perf_counter_ns() - result = func(*args, **kwargs) - end = time.perf_counter_ns() - - os.environ["CODEFLASH_BENCHMARKING"] = "False" - CodeFlashBenchmarkPlugin.benchmark_timings.append((benchmark_file_name, benchmark_function_name, line_number, end - start)) - return result - - return Benchmark() + return CodeFlashBenchmarkPlugin.Benchmark(request) \ No newline at end of file diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 685bfe739..becf606a4 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,17 +1,23 @@ -def print_benchmark_table(function_benchmark_timings: dict[str,dict[str,int]], total_benchmark_timings: dict[str,int]): - # Define column widths - benchmark_col_width = 50 - time_col_width = 15 +from rich.console import Console +from rich.table import Table - # Print table header - header = f"{'Benchmark Test':{benchmark_col_width}} | {'Total Time (ms)':{time_col_width}} | {'Function Time (ms)':{time_col_width}} | {'Percentage (%)':{time_col_width}}" - print(header) - print("-" * len(header)) + +def print_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]], + total_benchmark_timings: dict[str, int]): + console = Console() # Process each function's benchmark data for func_path, test_times in function_benchmark_timings.items(): function_name = func_path.split(":")[-1] - print(f"\n== Function: {function_name} ==") + + # 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) + 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") # Sort by percentage (highest first) sorted_tests = [] @@ -26,9 +32,16 @@ def print_benchmark_table(function_benchmark_timings: dict[str,dict[str,int]], t sorted_tests.sort(key=lambda x: x[3], reverse=True) - # Print each test's data + # Add rows to the table for test_name, total_time, func_time, percentage in sorted_tests: benchmark_file, benchmark_func, benchmark_line = test_name.split("::") benchmark_name = f"{benchmark_file}::{benchmark_func}" - print(f"{benchmark_name:{benchmark_col_width}} | {total_time:{time_col_width}.3f} | {func_time:{time_col_width}.3f} | {percentage:{time_col_width}.2f}") - print() + table.add_row( + benchmark_name, + f"{total_time:.3f}", + f"{func_time:.3f}", + f"{percentage:.2f}" + ) + + # Print the table + console.print(table) \ No newline at end of file From 82cb7a9e0e720b72872b4466a72b6ff8e1e154d4 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 15:11:17 -0700 Subject: [PATCH 072/122] cleanup tests better --- tests/test_trace_benchmarks.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index c49e7c693..4bc0ef278 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -150,5 +150,5 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): assert test_sort_path.read_text("utf-8").strip()==test_sort_code.strip() finally: # cleanup - shutil.rmtree(tests_root) - pass \ No newline at end of file + if tests_root.exists(): + shutil.rmtree(tests_root, ignore_errors=True) \ No newline at end of file From 7601895651a92156541456594433f19952c234ba Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 15:49:26 -0700 Subject: [PATCH 073/122] modified paths in test --- tests/test_trace_benchmarks.py | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 4bc0ef278..4113d954e 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -32,34 +32,36 @@ def test_trace_benchmarks(): # Assert the length of function calls assert len(function_calls) == 7, f"Expected 6 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() # Expected function calls expected_calls = [ ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_class_sort", "test_benchmark_bubble_sort.py", 20), ("sort_class", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_class_sort", "test_benchmark_bubble_sort.py", 18), ("sort_static", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_class_sort", "test_benchmark_bubble_sort.py", 19), ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_class_sort", "test_benchmark_bubble_sort.py", 17), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_sort", "test_benchmark_bubble_sort.py", 7), ("compute_and_sort", "", "code_to_optimize.process_and_bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/process_and_bubble_sort_codeflash_trace.py'}", + f"{process_and_bubble_sort_path}", "test_compute_and_sort", "test_process_and_sort.py", 4), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", - f"{project_root / 'code_to_optimize/bubble_sort_codeflash_trace.py'}", + f"{bubble_sort_path}", "test_no_func", "test_process_and_sort.py", 8), ] for idx, (actual, expected) in enumerate(zip(function_calls, expected_calls)): @@ -86,7 +88,7 @@ def test_trace_benchmarks(): trace_file_path = r"{output_file.as_posix()}" def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "sorter" @@ -99,7 +101,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): ret = instance.sorter(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) if not args: @@ -107,13 +109,13 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_class(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_static(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_static(*args, **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "__init__" @@ -141,7 +143,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): trace_file_path = r"{output_file}" def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"/Users/alvinryanputra/cf/codeflash/code_to_optimize/bubble_sort_codeflash_trace.py", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"{bubble_sort_path}", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_sorter(*args, **kwargs) @@ -150,5 +152,5 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): assert test_sort_path.read_text("utf-8").strip()==test_sort_code.strip() finally: # cleanup - if tests_root.exists(): - shutil.rmtree(tests_root, ignore_errors=True) \ No newline at end of file + shutil.rmtree(tests_root) + pass \ No newline at end of file From 4d6942708110161cd20baed8959a801ae589c0a1 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 15:53:17 -0700 Subject: [PATCH 074/122] typing fix --- codeflash/models/models.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 5358b8e4e..a91f3b42a 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -76,11 +76,11 @@ class BestOptimization(BaseModel): candidate: OptimizedCandidate helper_functions: list[FunctionSource] runtime: int - replay_runtime: int | None - replay_performance_gain: float | None + replay_runtime: Optional[int] = None + replay_performance_gain: Optional[float] = None winning_behavioral_test_results: TestResults winning_benchmarking_test_results: TestResults - winning_replay_benchmarking_test_results : TestResults | None = None + winning_replay_benchmarking_test_results : Optional[TestResults] = None class CodeString(BaseModel): @@ -227,7 +227,7 @@ class OriginalCodeBaseline(BaseModel): benchmarking_test_results: TestResults line_profile_results: dict runtime: int - coverage_results: CoverageData | None + coverage_results: Optional[CoverageData] class CoverageStatus(Enum): From ebe3e126afe33be437f1ac8e32aca764d44574d5 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 20 Mar 2025 17:01:29 -0700 Subject: [PATCH 075/122] typing fix for 3.9 --- codeflash/result/explanation.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 77528a8a6..8aec614cc 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -1,4 +1,6 @@ +from __future__ import annotations from pathlib import Path +from typing import Optional, Union from pydantic.dataclasses import dataclass @@ -15,9 +17,9 @@ class Explanation: best_runtime_ns: int function_name: str file_path: Path - replay_performance_gain: float | None - fto_benchmark_timings: dict[str, int] | None - total_benchmark_timings: dict[str, int] | None + replay_performance_gain: Optional[float] + fto_benchmark_timings: Optional[Union[dict, int]] + total_benchmark_timings: Optional[Union[dict, int]] @property def perf_improvement_line(self) -> str: From 0449d0d4cff193b4817435954f83ef91131a8121 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 21 Mar 2025 13:09:28 -0700 Subject: [PATCH 076/122] typing fix for 3.9 --- codeflash/result/explanation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 8aec614cc..60f0d9a70 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -18,8 +18,8 @@ class Explanation: function_name: str file_path: Path replay_performance_gain: Optional[float] - fto_benchmark_timings: Optional[Union[dict, int]] - total_benchmark_timings: Optional[Union[dict, int]] + fto_benchmark_timings: Optional[dict[str, int]] + total_benchmark_timings: Optional[dict[str, int]] @property def perf_improvement_line(self) -> str: From baac96451eba651db6a43e09122013df23f9b257 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Mon, 24 Mar 2025 16:45:13 -0700 Subject: [PATCH 077/122] works with multithreading, added test --- code_to_optimize/bubble_sort_multithread.py | 23 +++++++ .../benchmarks/test_benchmark_bubble_sort.py | 9 +-- .../benchmarks/test_process_and_sort.py | 4 +- .../test_multithread_sort.py | 4 ++ .../test_benchmark_bubble_sort.py | 20 ++++++ .../benchmarks_test/test_process_and_sort.py | 8 +++ codeflash/benchmarking/codeflash_trace.py | 16 ++--- codeflash/benchmarking/utils.py | 61 ++++++++++++------- codeflash/optimization/function_optimizer.py | 4 +- codeflash/optimization/optimizer.py | 16 ++--- codeflash/result/explanation.py | 22 ++++--- tests/test_trace_benchmarks.py | 61 ++++++++++++++++++- 12 files changed, 185 insertions(+), 63 deletions(-) create mode 100644 code_to_optimize/bubble_sort_multithread.py create mode 100644 code_to_optimize/tests/pytest/benchmarks_multithread/test_multithread_sort.py create mode 100644 code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py create mode 100644 code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py diff --git a/code_to_optimize/bubble_sort_multithread.py b/code_to_optimize/bubble_sort_multithread.py new file mode 100644 index 000000000..3659b01bf --- /dev/null +++ b/code_to_optimize/bubble_sort_multithread.py @@ -0,0 +1,23 @@ +# from code_to_optimize.bubble_sort_codeflash_trace import sorter +from code_to_optimize.bubble_sort_codeflash_trace import sorter +import concurrent.futures + + +def multithreaded_sorter(unsorted_lists: list[list[int]]) -> list[list[int]]: + # Create a list to store results in the correct order + sorted_lists = [None] * len(unsorted_lists) + + # Use ThreadPoolExecutor to manage threads + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + # Submit all sorting tasks and map them to their original indices + future_to_index = { + executor.submit(sorter, unsorted_list): i + for i, unsorted_list in enumerate(unsorted_lists) + } + + # Collect results as they complete + for future in concurrent.futures.as_completed(future_to_index): + index = future_to_index[future] + sorted_lists[index] = future.result() + + return sorted_lists \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py index 03b9d38d1..3d7b24a6c 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_benchmark_bubble_sort.py @@ -1,6 +1,6 @@ import pytest -from code_to_optimize.bubble_sort_codeflash_trace import sorter, Sorter +from code_to_optimize.bubble_sort import sorter def test_sort(benchmark): @@ -11,10 +11,3 @@ def test_sort(benchmark): def test_sort2(): result = sorter(list(reversed(range(500)))) assert result == list(range(500)) - -def test_class_sort(benchmark): - obj = Sorter(list(reversed(range(100)))) - result1 = benchmark(obj.sorter, 2) - result2 = benchmark(Sorter.sort_class, list(reversed(range(100)))) - result3 = benchmark(Sorter.sort_static, list(reversed(range(100)))) - result4 = benchmark(Sorter, [1,2,3]) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py index bcd42eab9..8d31c926a 100644 --- a/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks/test_process_and_sort.py @@ -1,5 +1,5 @@ -from code_to_optimize.process_and_bubble_sort_codeflash_trace import compute_and_sort -from code_to_optimize.bubble_sort_codeflash_trace import sorter +from code_to_optimize.process_and_bubble_sort import compute_and_sort +from code_to_optimize.bubble_sort import sorter def test_compute_and_sort(benchmark): result = benchmark(compute_and_sort, list(reversed(range(500)))) assert result == 62208.5 diff --git a/code_to_optimize/tests/pytest/benchmarks_multithread/test_multithread_sort.py b/code_to_optimize/tests/pytest/benchmarks_multithread/test_multithread_sort.py new file mode 100644 index 000000000..4a5c68a2b --- /dev/null +++ b/code_to_optimize/tests/pytest/benchmarks_multithread/test_multithread_sort.py @@ -0,0 +1,4 @@ +from code_to_optimize.bubble_sort_multithread import multithreaded_sorter + +def test_benchmark_sort(benchmark): + benchmark(multithreaded_sorter, [list(range(1000)) for i in range (10)]) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py new file mode 100644 index 000000000..03b9d38d1 --- /dev/null +++ b/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py @@ -0,0 +1,20 @@ +import pytest + +from code_to_optimize.bubble_sort_codeflash_trace import sorter, Sorter + + +def test_sort(benchmark): + result = benchmark(sorter, list(reversed(range(500)))) + assert result == list(range(500)) + +# This should not be picked up as a benchmark test +def test_sort2(): + result = sorter(list(reversed(range(500)))) + assert result == list(range(500)) + +def test_class_sort(benchmark): + obj = Sorter(list(reversed(range(100)))) + result1 = benchmark(obj.sorter, 2) + result2 = benchmark(Sorter.sort_class, list(reversed(range(100)))) + result3 = benchmark(Sorter.sort_static, list(reversed(range(100)))) + result4 = benchmark(Sorter, [1,2,3]) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py new file mode 100644 index 000000000..bcd42eab9 --- /dev/null +++ b/code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py @@ -0,0 +1,8 @@ +from code_to_optimize.process_and_bubble_sort_codeflash_trace import compute_and_sort +from code_to_optimize.bubble_sort_codeflash_trace import sorter +def test_compute_and_sort(benchmark): + result = benchmark(compute_and_sort, list(reversed(range(500)))) + assert result == 62208.5 + +def test_no_func(benchmark): + benchmark(sorter, list(reversed(range(500)))) \ No newline at end of file diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index b1236ffbf..f708d752f 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -15,11 +15,6 @@ class CodeflashTrace: def __init__(self) -> None: self.function_calls_data = [] - # def __enter__(self) -> None: - # # Initialize for context manager use - # self.function_calls_data = [] - # return self - def __exit__(self, exc_type, exc_val, exc_tb) -> None: # Cleanup is optional here pass @@ -37,15 +32,14 @@ def __call__(self, func: Callable) -> Callable: @functools.wraps(func) def wrapper(*args, **kwargs): # Measure execution time - start_time = time.perf_counter_ns() + start_time = time.thread_time_ns() result = func(*args, **kwargs) - end_time = time.perf_counter_ns() - + end_time = time.thread_time_ns() # Calculate execution time execution_time = end_time - start_time # Measure overhead - overhead_start_time = time.perf_counter_ns() + overhead_start_time = time.thread_time_ns() try: # Check if currently in pytest benchmark fixture @@ -66,7 +60,7 @@ def wrapper(*args, **kwargs): if "." in qualname: class_name = qualname.split(".")[0] # Calculate overhead time - overhead_end_time = time.perf_counter_ns() + overhead_end_time = time.thread_time_ns() overhead_time = overhead_end_time - overhead_start_time @@ -75,7 +69,7 @@ def wrapper(*args, **kwargs): benchmark_function_name, benchmark_file_name, benchmark_line_number, execution_time, overhead_time, pickled_args, pickled_kwargs) ) - + print("appended") except Exception as e: print(f"Error in codeflash_trace: {e}") diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index becf606a4..eeacb6975 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,47 +1,64 @@ from rich.console import Console from rich.table import Table +from codeflash.cli_cmds.console import logger -def print_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]], - total_benchmark_timings: dict[str, int]): - console = Console() +def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]], + total_benchmark_timings: dict[str, int]) -> dict[str, list[tuple[str, float, float, float]]]: + function_to_result = {} # Process each function's benchmark data for func_path, test_times in function_benchmark_timings.items(): - 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) - 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") - # Sort by percentage (highest first) sorted_tests = [] for test_name, func_time in test_times.items(): total_time = total_benchmark_timings.get(test_name, 0) + if func_time > total_time: + logger.debug(f"Skipping test {test_name} due to func_time {func_time} > total_time {total_time}") + # If the function time is greater than total time, likely to have multithreading / multiprocessing issues. + # Do not try to project the optimization impact for this function. + sorted_tests.append((test_name, 0.0, 0.0, 0.0)) if total_time > 0: percentage = (func_time / total_time) * 100 # Convert nanoseconds to milliseconds func_time_ms = func_time / 1_000_000 total_time_ms = total_time / 1_000_000 sorted_tests.append((test_name, total_time_ms, func_time_ms, percentage)) - sorted_tests.sort(key=lambda x: x[3], reverse=True) + function_to_result[func_path] = sorted_tests + return function_to_result + +def print_benchmark_table(function_to_results: dict[str, list[tuple[str, float, float, float]]]) -> None: + console = Console() + for func_path, sorted_tests in function_to_results.items(): + 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) + 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") - # Add rows to the table for test_name, total_time, func_time, percentage in sorted_tests: benchmark_file, benchmark_func, benchmark_line = test_name.split("::") benchmark_name = f"{benchmark_file}::{benchmark_func}" - table.add_row( - benchmark_name, - f"{total_time:.3f}", - f"{func_time:.3f}", - f"{percentage:.2f}" - ) + if total_time == 0.0: + table.add_row( + benchmark_name, + "N/A", + "N/A", + "N/A" + ) + else: + table.add_row( + benchmark_name, + f"{total_time:.3f}", + f"{func_time:.3f}", + f"{percentage:.2f}" + ) # Print the table console.print(table) \ No newline at end of file diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index b2af5e0bf..953a12028 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -92,7 +92,7 @@ def __init__( function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_ast: ast.FunctionDef | None = None, aiservice_client: AiServiceClient | None = None, - function_benchmark_timings: dict[str, dict[str, int]] | None = None, + function_benchmark_timings: dict[str, int] | None = None, total_benchmark_timings: dict[str, int] | None = None, args: Namespace | None = None, ) -> None: @@ -277,7 +277,7 @@ def optimize_function(self) -> Result[BestOptimization, str]: function_name=function_to_optimize_qualified_name, file_path=self.function_to_optimize.file_path, replay_performance_gain=best_optimization.replay_performance_gain if self.args.benchmark else None, - fto_benchmark_timings = self.function_benchmark_timings[self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root)] if self.args.benchmark else None, + fto_benchmark_timings = self.function_benchmark_timings if self.args.benchmark else None, total_benchmark_timings = self.total_benchmark_timings if self.args.benchmark else None, ) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 894a8911e..d0152b849 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -10,10 +10,9 @@ from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest -from codeflash.benchmarking.utils import print_benchmark_table +from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table from codeflash.cli_cmds.console import console, logger, progress_bar from codeflash.code_utils import env_utils -from codeflash.code_utils.code_extractor import add_needed_imports_from_module from codeflash.code_utils.code_replacer import normalize_code, normalize_node from codeflash.code_utils.code_utils import get_run_tmp_file from codeflash.code_utils.static_analysis import analyze_imported_modules, get_first_top_level_function_or_method_ast @@ -115,15 +114,15 @@ def run(self) -> None: instrument_codeflash_trace_decorator(fto) trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" replay_tests_dir = Path(self.args.tests_root) / "codeflash_replay_tests" - trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Simply run all tests that use pytest-benchmark + trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Run all tests that use pytest-benchmark replay_count = generate_replay_test(trace_file, replay_tests_dir) if replay_count == 0: logger.info(f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization") else: function_benchmark_timings = get_function_benchmark_timings(trace_file) total_benchmark_timings = get_benchmark_timings(trace_file) - - print_benchmark_table(function_benchmark_timings, total_benchmark_timings) + function_to_results = validate_and_format_benchmark_table(function_benchmark_timings, total_benchmark_timings) + print_benchmark_table(function_to_results) logger.info("Finished tracing existing benchmarks") except Exception as e: logger.info(f"Error while tracing existing benchmarks: {e}") @@ -212,9 +211,12 @@ def run(self) -> None: f"Skipping optimization." ) continue - if self.args.benchmark and function_benchmark_timings and total_benchmark_timings: + qualified_name_w_module = function_to_optimize.qualified_name_with_modules_from_root( + self.args.project_root + ) + if self.args.benchmark and function_benchmark_timings and qualified_name_w_module in function_benchmark_timings and total_benchmark_timings: function_optimizer = self.create_function_optimizer( - function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code, function_benchmark_timings, total_benchmark_timings + function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_code[original_module_path].source_code, function_benchmark_timings[qualified_name_w_module], total_benchmark_timings ) else: function_optimizer = self.create_function_optimizer( diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 60f0d9a70..c44282e88 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -54,16 +54,18 @@ def to_console_string(self) -> str: continue total_benchmark_timing = self.total_benchmark_timings[benchmark_key] - # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values - benchmark_info += f"Original timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(total_benchmark_timing)}\n" - replay_speedup = self.replay_performance_gain - expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / ( - replay_speedup + 1) * og_benchmark_timing - benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(int(expected_new_benchmark_timing))}\n" - - benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing - benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 - benchmark_info += f"Benchmark speedup for {benchmark_file_name}::{benchmark_test_function}: {benchmark_speedup_percent:.2f}%\n\n" + if total_benchmark_timing == 0: + benchmark_info += f"Benchmark timing for {benchmark_file_name}::{benchmark_test_function} was improved, but the speedup cannot be estimated.\n" + else: + # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values + benchmark_info += f"Original timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(total_benchmark_timing)}\n" + replay_speedup = self.replay_performance_gain + expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / ( + replay_speedup + 1) * og_benchmark_timing + benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(int(expected_new_benchmark_timing))}\n" + benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing + benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 + benchmark_info += f"Benchmark speedup for {benchmark_file_name}::{benchmark_test_function}: {benchmark_speedup_percent:.2f}%\n\n" return ( f"Optimized {self.function_name} in {self.file_path}\n" diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 4113d954e..244b08029 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,9 +1,12 @@ import sqlite3 from codeflash.benchmarking.codeflash_trace import codeflash_trace +from codeflash.benchmarking.get_trace_info import get_function_benchmark_timings, get_benchmark_timings from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.replay_test import generate_replay_test from pathlib import Path + +from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table from codeflash.code_utils.code_utils import get_run_tmp_file import shutil @@ -11,7 +14,7 @@ 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" + benchmarks_root = project_root / "tests" / "pytest" / "benchmarks_test" tests_root = project_root / "tests" / "test_trace_benchmarks" tests_root.mkdir(parents=False, exist_ok=False) output_file = (tests_root / Path("test_trace_benchmarks.trace")).resolve() @@ -150,6 +153,62 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): """ assert test_sort_path.read_text("utf-8").strip()==test_sort_code.strip() + finally: + # cleanup + shutil.rmtree(tests_root) + pass + +def test_trace_multithreaded_benchmark() -> None: + project_root = Path(__file__).parent.parent / "code_to_optimize" + benchmarks_root = project_root / "tests" / "pytest" / "benchmarks_multithread" + tests_root = project_root / "tests" / "test_trace_benchmarks" + tests_root.mkdir(parents=False, exist_ok=False) + output_file = (tests_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_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM function_calls ORDER BY benchmark_file_name, benchmark_function_name, function_name") + function_calls = cursor.fetchall() + + # Assert the length of function calls + assert len(function_calls) == 10, f"Expected 10 function calls, but got {len(function_calls)}" + function_benchmark_timings = get_function_benchmark_timings(output_file) + total_benchmark_timings = get_benchmark_timings(output_file) + # This will throw an error if summed function timings exceed total benchmark timing + 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", "test_multithread_sort.py", 4), + ] + 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_name" + 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_file_name" + assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number" + # Close connection + conn.close() + finally: # cleanup shutil.rmtree(tests_root) From 357f5863b4916579004138eff9aa70b37b72b2f3 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 25 Mar 2025 09:51:12 -0700 Subject: [PATCH 078/122] refactored get_function_benchmark_timings and get_benchmark_timings into BenchmarkDatabaseUtils class --- .../benchmarking/benchmark_database_utils.py | 117 +++++++++++++++++ codeflash/benchmarking/get_trace_info.py | 121 ------------------ codeflash/optimization/optimizer.py | 6 +- tests/test_trace_benchmarks.py | 9 +- 4 files changed, 123 insertions(+), 130 deletions(-) delete mode 100644 codeflash/benchmarking/get_trace_info.py diff --git a/codeflash/benchmarking/benchmark_database_utils.py b/codeflash/benchmarking/benchmark_database_utils.py index b9b36079d..1c117553c 100644 --- a/codeflash/benchmarking/benchmark_database_utils.py +++ b/codeflash/benchmarking/benchmark_database_utils.py @@ -177,3 +177,120 @@ def close(self) -> None: self.connection.close() self.connection = None + + @staticmethod + def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, int]]: + """Process the trace file and extract timing data for all functions. + + Args: + trace_path: Path to the trace file + + Returns: + A nested dictionary where: + - Outer keys are module_name.qualified_name (module.class.function) + - Inner keys are benchmark filename :: benchmark test function :: line number + - Values are function timing in milliseconds + + """ + # Initialize the result dictionary + result = {} + + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + + try: + # Query the function_calls table for all function calls + cursor.execute( + "SELECT module_name, class_name, function_name, " + "benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns " + "FROM function_calls" + ) + + # Process each row + for row in cursor.fetchall(): + module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row + + # Create the function key (module_name.class_name.function_name) + if class_name: + qualified_name = f"{module_name}.{class_name}.{function_name}" + else: + qualified_name = f"{module_name}.{function_name}" + + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + + # Initialize the inner dictionary if needed + if qualified_name not in result: + result[qualified_name] = {} + + # If multiple calls to the same function in the same benchmark, + # add the times together + if benchmark_key in result[qualified_name]: + result[qualified_name][benchmark_key] += time_ns + else: + result[qualified_name][benchmark_key] = time_ns + + finally: + # Close the connection + connection.close() + + return result + + @staticmethod + def get_benchmark_timings(trace_path: Path) -> dict[str, int]: + """Extract total benchmark timings from trace files. + + Args: + trace_path: Path to the trace file + + Returns: + A dictionary mapping where: + - Keys are benchmark filename :: benchmark test function :: line number + - Values are total benchmark timing in milliseconds (with overhead subtracted) + + """ + # Initialize the result dictionary + result = {} + overhead_by_benchmark = {} + + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + + try: + # Query the function_calls table to get total overhead for each benchmark + cursor.execute( + "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " + "FROM function_calls " + "GROUP BY benchmark_file_name, 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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + 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_name, benchmark_function_name, benchmark_line_number, time_ns " + "FROM benchmark_timings" + ) + + # Process each row and subtract overhead + for row in cursor.fetchall(): + benchmark_file, benchmark_func, benchmark_line, time_ns = row + + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + + # Subtract overhead from total time + overhead = overhead_by_benchmark.get(benchmark_key, 0) + result[benchmark_key] = time_ns - overhead + + finally: + # Close the connection + connection.close() + + return result diff --git a/codeflash/benchmarking/get_trace_info.py b/codeflash/benchmarking/get_trace_info.py deleted file mode 100644 index d43327af7..000000000 --- a/codeflash/benchmarking/get_trace_info.py +++ /dev/null @@ -1,121 +0,0 @@ -import sqlite3 -from pathlib import Path - -from codeflash.discovery.functions_to_optimize import FunctionToOptimize - - -def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, int]]: - """Process the trace file and extract timing data for all functions. - - Args: - trace_path: Path to the trace file - - Returns: - A nested dictionary where: - - Outer keys are module_name.qualified_name (module.class.function) - - Inner keys are benchmark filename :: benchmark test function :: line number - - Values are function timing in milliseconds - - """ - # Initialize the result dictionary - result = {} - - # Connect to the SQLite database - connection = sqlite3.connect(trace_path) - cursor = connection.cursor() - - try: - # Query the function_calls table for all function calls - cursor.execute( - "SELECT module_name, class_name, function_name, " - "benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns " - "FROM function_calls" - ) - - # Process each row - for row in cursor.fetchall(): - module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row - - # Create the function key (module_name.class_name.function_name) - if class_name: - qualified_name = f"{module_name}.{class_name}.{function_name}" - else: - qualified_name = f"{module_name}.{function_name}" - - # Create the benchmark key (file::function::line) - benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - - # Initialize the inner dictionary if needed - if qualified_name not in result: - result[qualified_name] = {} - - # If multiple calls to the same function in the same benchmark, - # add the times together - if benchmark_key in result[qualified_name]: - result[qualified_name][benchmark_key] += time_ns - else: - result[qualified_name][benchmark_key] = time_ns - - finally: - # Close the connection - connection.close() - - return result - - -def get_benchmark_timings(trace_path: Path) -> dict[str, int]: - """Extract total benchmark timings from trace files. - - Args: - trace_path: Path to the trace file - - Returns: - A dictionary mapping where: - - Keys are benchmark filename :: benchmark test function :: line number - - Values are total benchmark timing in milliseconds (with overhead subtracted) - - """ - # Initialize the result dictionary - result = {} - overhead_by_benchmark = {} - - # Connect to the SQLite database - connection = sqlite3.connect(trace_path) - cursor = connection.cursor() - - try: - # Query the function_calls table to get total overhead for each benchmark - cursor.execute( - "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " - "FROM function_calls " - "GROUP BY benchmark_file_name, 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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - 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_name, benchmark_function_name, benchmark_line_number, time_ns " - "FROM benchmark_timings" - ) - - # Process each row and subtract overhead - for row in cursor.fetchall(): - benchmark_file, benchmark_func, benchmark_line, time_ns = row - - # Create the benchmark key (file::function::line) - benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - - # Subtract overhead from total time - overhead = overhead_by_benchmark.get(benchmark_key, 0) - result[benchmark_key] = time_ns - overhead - - finally: - # Close the connection - connection.close() - - return result diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index d0152b849..08db413b2 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -8,6 +8,7 @@ from typing import TYPE_CHECKING from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient +from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table @@ -24,7 +25,6 @@ from codeflash.telemetry.posthog_cf import ph from codeflash.verification.test_results import TestType 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 codeflash.benchmarking.instrument_codeflash_trace import instrument_codeflash_trace_decorator @@ -119,8 +119,8 @@ def run(self) -> None: if replay_count == 0: logger.info(f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization") else: - function_benchmark_timings = get_function_benchmark_timings(trace_file) - total_benchmark_timings = get_benchmark_timings(trace_file) + function_benchmark_timings = BenchmarkDatabaseUtils.get_function_benchmark_timings(trace_file) + total_benchmark_timings = BenchmarkDatabaseUtils.get_benchmark_timings(trace_file) function_to_results = validate_and_format_benchmark_table(function_benchmark_timings, total_benchmark_timings) print_benchmark_table(function_to_results) logger.info("Finished tracing existing benchmarks") diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 244b08029..fcc5b0f67 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,13 +1,11 @@ import sqlite3 -from codeflash.benchmarking.codeflash_trace import codeflash_trace -from codeflash.benchmarking.get_trace_info import get_function_benchmark_timings, get_benchmark_timings +from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.replay_test import generate_replay_test from pathlib import Path from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table -from codeflash.code_utils.code_utils import get_run_tmp_file import shutil @@ -180,9 +178,8 @@ def test_trace_multithreaded_benchmark() -> None: # Assert the length of function calls assert len(function_calls) == 10, f"Expected 10 function calls, but got {len(function_calls)}" - function_benchmark_timings = get_function_benchmark_timings(output_file) - total_benchmark_timings = get_benchmark_timings(output_file) - # This will throw an error if summed function timings exceed total benchmark timing + function_benchmark_timings = BenchmarkDatabaseUtils.get_function_benchmark_timings(output_file) + total_benchmark_timings = BenchmarkDatabaseUtils.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 From 9efa47f942c88bc4378cf27f13147f90b7376353 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 25 Mar 2025 10:28:09 -0700 Subject: [PATCH 079/122] fixed isort --- codeflash/benchmarking/instrument_codeflash_trace.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py index 93f51baed..017cecaca 100644 --- a/codeflash/benchmarking/instrument_codeflash_trace.py +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -98,12 +98,12 @@ def instrument_codeflash_trace_decorator( """Instrument __init__ function with codeflash_trace decorator if it's in a class.""" # Instrument fto class original_code = function_to_optimize.file_path.read_text(encoding="utf-8") - - # Modify the code - modified_code = isort.code(add_codeflash_decorator_to_code( + new_code = add_codeflash_decorator_to_code( original_code, function_to_optimize - )) + ) + # Modify the code + modified_code = isort.code(code=new_code, float_to_top=True) # Write the modified code back to the file function_to_optimize.file_path.write_text(modified_code, encoding="utf-8") From 64b4c6406fb9ed32e785bbbee56d4ac0ad0c8767 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 26 Mar 2025 10:20:28 -0700 Subject: [PATCH 080/122] modified PR info --- codeflash/benchmarking/utils.py | 61 ++- codeflash/github/PrComment.py | 17 +- codeflash/models/models.py | 521 ++++++++++--------- codeflash/optimization/function_optimizer.py | 12 +- codeflash/result/create_pr.py | 2 + codeflash/result/explanation.py | 35 +- 6 files changed, 355 insertions(+), 293 deletions(-) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index eeacb6975..38c31b55b 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,7 +1,12 @@ +from __future__ import annotations +from typing import Optional + from rich.console import Console from rich.table import Table from codeflash.cli_cmds.console import logger +from codeflash.code_utils.time_utils import humanize_runtime +from codeflash.models.models import ProcessedBenchmarkInfo, BenchmarkDetail def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]], @@ -61,4 +66,58 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[str, float, ) # Print the table - console.print(table) \ No newline at end of file + console.print(table) + + +def process_benchmark_data( + replay_performance_gain: float, + fto_benchmark_timings: dict[str, int], + total_benchmark_timings: dict[str, int] +) -> Optional[ProcessedBenchmarkInfo]: + """Process benchmark data and generate detailed benchmark information. + + Args: + replay_performance_gain: The performance gain from replay + fto_benchmark_timings: Function to optimize benchmark timings + total_benchmark_timings: Total benchmark timings + + Returns: + ProcessedBenchmarkInfo containing processed benchmark details + + """ + if not replay_performance_gain or not fto_benchmark_timings or not total_benchmark_timings: + return None + + benchmark_details = [] + + for benchmark_key, og_benchmark_timing in fto_benchmark_timings.items(): + try: + benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") + except ValueError: + continue # Skip malformed benchmark keys + + total_benchmark_timing = total_benchmark_timings.get(benchmark_key, 0) + + if total_benchmark_timing == 0: + continue # Skip benchmarks with zero timing + + # Calculate expected new benchmark timing + expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + ( + 1 / (replay_performance_gain + 1) + ) * og_benchmark_timing + + # Calculate speedup + benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing + benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 + + benchmark_details.append( + BenchmarkDetail( + benchmark_name=benchmark_file_name, + test_function=benchmark_test_function, + original_timing=humanize_runtime(int(total_benchmark_timing)), + expected_new_timing=humanize_runtime(int(expected_new_benchmark_timing)), + speedup_percent=benchmark_speedup_percent + ) + ) + + return ProcessedBenchmarkInfo(benchmark_details=benchmark_details) \ No newline at end of file diff --git a/codeflash/github/PrComment.py b/codeflash/github/PrComment.py index d7c12d962..4ef162cda 100644 --- a/codeflash/github/PrComment.py +++ b/codeflash/github/PrComment.py @@ -1,10 +1,12 @@ -from typing import Union +from __future__ import annotations +from typing import Union, Optional from pydantic import BaseModel from pydantic.dataclasses import dataclass from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.models.models import TestResults +from codeflash.models.models import BenchmarkDetail +from codeflash.verification.test_results import TestResults @dataclass(frozen=True, config={"arbitrary_types_allowed": True}) @@ -18,15 +20,16 @@ class PrComment: speedup_pct: str winning_behavioral_test_results: TestResults winning_benchmarking_test_results: TestResults + benchmark_details: Optional[list[BenchmarkDetail]] = None - def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str]]: + def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str, Optional[list[dict[str, any]]]]]: report_table = { test_type.to_name(): result for test_type, result in self.winning_behavioral_test_results.get_test_pass_fail_report_by_type().items() if test_type.to_name() } - return { + result = { "optimization_explanation": self.optimization_explanation, "best_runtime": humanize_runtime(self.best_runtime), "original_runtime": humanize_runtime(self.original_runtime), @@ -38,6 +41,12 @@ def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str]]: "report_table": report_table, } + # Add benchmark details if available + if self.benchmark_details: + result["benchmark_details"] = self.benchmark_details + + return result + class FileDiffContent(BaseModel): oldContent: str diff --git a/codeflash/models/models.py b/codeflash/models/models.py index a91f3b42a..52d1e4285 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -1,30 +1,30 @@ from __future__ import annotations -from typing import TYPE_CHECKING - -from rich.tree import Tree - -from codeflash.cli_cmds.console import DEBUG_MODE - -if TYPE_CHECKING: - from collections.abc import Iterator import enum +import json import re -import sys from collections.abc import Collection, Iterator from enum import Enum, IntEnum from pathlib import Path from re import Pattern -from typing import Annotated, Optional, cast +from typing import Annotated, Any, Optional, Union +import sentry_sdk +from coverage.exceptions import NoDataError from jedi.api.classes import Name from pydantic import AfterValidator, BaseModel, ConfigDict, Field from pydantic.dataclasses import dataclass from codeflash.cli_cmds.console import console, logger from codeflash.code_utils.code_utils import validate_python_code +from codeflash.code_utils.coverage_utils import ( + build_fully_qualified_name, + extract_dependent_function, + generate_candidates, +) from codeflash.code_utils.env_utils import is_end_to_end -from codeflash.verification.comparator import comparator +from codeflash.code_utils.time_utils import humanize_runtime +from codeflash.verification.test_results import TestResults, TestType # If the method spam is in the class Ham, which is at the top level of the module eggs in the package foo, the fully # qualified name of the method is foo.eggs.Ham.spam, its qualified name is Ham.spam, and its name is spam. The full name @@ -58,19 +58,15 @@ class FunctionSource: def __eq__(self, other: object) -> bool: if not isinstance(other, FunctionSource): return False - return ( - self.file_path == other.file_path - and self.qualified_name == other.qualified_name - and self.fully_qualified_name == other.fully_qualified_name - and self.only_function_name == other.only_function_name - and self.source_code == other.source_code - ) + return (self.file_path == other.file_path and + self.qualified_name == other.qualified_name and + self.fully_qualified_name == other.fully_qualified_name and + self.only_function_name == other.only_function_name and + self.source_code == other.source_code) def __hash__(self) -> int: - return hash( - (self.file_path, self.qualified_name, self.fully_qualified_name, self.only_function_name, self.source_code) - ) - + return hash((self.file_path, self.qualified_name, self.fully_qualified_name, + self.only_function_name, self.source_code)) class BestOptimization(BaseModel): candidate: OptimizedCandidate @@ -82,7 +78,47 @@ class BestOptimization(BaseModel): winning_benchmarking_test_results: TestResults winning_replay_benchmarking_test_results : Optional[TestResults] = None +@dataclass +class BenchmarkDetail: + benchmark_name: str + test_function: str + original_timing: str + expected_new_timing: str + speedup_percent: float + + def to_string(self) -> str: + return ( + f"Original timing for {self.benchmark_name}::{self.test_function}: {self.original_timing}\n" + f"Expected new timing for {self.benchmark_name}::{self.test_function}: {self.expected_new_timing}\n" + f"Benchmark speedup for {self.benchmark_name}::{self.test_function}: {self.speedup_percent:.2f}%\n" + ) + + def to_dict(self) -> dict[str, any]: + return { + "benchmark_name": self.benchmark_name, + "test_function": self.test_function, + "original_timing": self.original_timing, + "expected_new_timing": self.expected_new_timing, + "speedup_percent": self.speedup_percent + } +@dataclass +class ProcessedBenchmarkInfo: + benchmark_details: list[BenchmarkDetail] + + def to_string(self) -> str: + if not self.benchmark_details: + return "" + + result = "Benchmark Performance Details:\n" + for detail in self.benchmark_details: + result += detail.to_string() + "\n" + return result + + def to_dict(self) -> dict[str, list[dict[str, any]]]: + return { + "benchmark_details": [detail.to_dict() for detail in self.benchmark_details] + } class CodeString(BaseModel): code: Annotated[str, AfterValidator(validate_python_code)] file_path: Optional[Path] = None @@ -107,8 +143,7 @@ class CodeOptimizationContext(BaseModel): read_writable_code: str = Field(min_length=1) read_only_context_code: str = "" helper_functions: list[FunctionSource] - preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]] - + preexisting_objects: set[tuple[str, tuple[FunctionParent,...]]] class CodeContextType(str, Enum): READ_WRITABLE = "READ_WRITABLE" @@ -225,7 +260,6 @@ class FunctionParent: class OriginalCodeBaseline(BaseModel): behavioral_test_results: TestResults benchmarking_test_results: TestResults - line_profile_results: dict runtime: int coverage_results: Optional[CoverageData] @@ -251,6 +285,209 @@ class CoverageData: blank_re: Pattern[str] = re.compile(r"\s*(#|$)") else_re: Pattern[str] = re.compile(r"\s*else\s*:\s*(#|$)") + @staticmethod + def load_from_sqlite_database( + database_path: Path, config_path: Path, function_name: str, code_context: CodeOptimizationContext, source_code_path: Path + ) -> CoverageData: + """Load coverage data from an SQLite database, mimicking the behavior of load_from_coverage_file.""" + from coverage import Coverage + from coverage.jsonreport import JsonReporter + + cov = Coverage(data_file=database_path,config_file=config_path, data_suffix=True, auto_data=True, branch=True) + + if not database_path.stat().st_size or not database_path.exists(): + logger.debug(f"Coverage database {database_path} is empty or does not exist") + sentry_sdk.capture_message(f"Coverage database {database_path} is empty or does not exist") + return CoverageData.create_empty(source_code_path, function_name, code_context) + cov.load() + + reporter = JsonReporter(cov) + temp_json_file = database_path.with_suffix(".report.json") + with temp_json_file.open("w") as f: + try: + reporter.report(morfs=[source_code_path.as_posix()], outfile=f) + except NoDataError: + sentry_sdk.capture_message(f"No coverage data found for {function_name} in {source_code_path}") + return CoverageData.create_empty(source_code_path, function_name, code_context) + with temp_json_file.open() as f: + original_coverage_data = json.load(f) + + coverage_data, status = CoverageData._parse_coverage_file(temp_json_file, source_code_path) + + main_func_coverage, dependent_func_coverage = CoverageData._fetch_function_coverages( + function_name, code_context, coverage_data, original_cov_data=original_coverage_data + ) + + total_executed_lines, total_unexecuted_lines = CoverageData._aggregate_coverage( + main_func_coverage, dependent_func_coverage + ) + + total_lines = total_executed_lines | total_unexecuted_lines + coverage = len(total_executed_lines) / len(total_lines) * 100 if total_lines else 0.0 + # coverage = (lines covered of the original function + its 1 level deep helpers) / (lines spanned by original function + its 1 level deep helpers), if no helpers then just the original function coverage + + functions_being_tested = [main_func_coverage.name] + if dependent_func_coverage: + functions_being_tested.append(dependent_func_coverage.name) + + graph = CoverageData._build_graph(main_func_coverage, dependent_func_coverage) + temp_json_file.unlink() + + return CoverageData( + file_path=source_code_path, + coverage=coverage, + function_name=function_name, + functions_being_tested=functions_being_tested, + graph=graph, + code_context=code_context, + main_func_coverage=main_func_coverage, + dependent_func_coverage=dependent_func_coverage, + status=status, + ) + + @staticmethod + def _parse_coverage_file( + coverage_file_path: Path, source_code_path: Path + ) -> tuple[dict[str, dict[str, Any]], CoverageStatus]: + with coverage_file_path.open() as f: + coverage_data = json.load(f) + + candidates = generate_candidates(source_code_path) + + logger.debug(f"Looking for coverage data in {' -> '.join(candidates)}") + for candidate in candidates: + try: + cov: dict[str, dict[str, Any]] = coverage_data["files"][candidate]["functions"] + logger.debug(f"Coverage data found for {source_code_path} in {candidate}") + status = CoverageStatus.PARSED_SUCCESSFULLY + break + except KeyError: + continue + else: + logger.debug(f"No coverage data found for {source_code_path} in {candidates}") + cov = {} + status = CoverageStatus.NOT_FOUND + return cov, status + + @staticmethod + def _fetch_function_coverages( + function_name: str, + code_context: CodeOptimizationContext, + coverage_data: dict[str, dict[str, Any]], + original_cov_data: dict[str, dict[str, Any]], + ) -> tuple[FunctionCoverage, Union[FunctionCoverage, None]]: + resolved_name = build_fully_qualified_name(function_name, code_context) + try: + main_function_coverage = FunctionCoverage( + name=resolved_name, + coverage=coverage_data[resolved_name]["summary"]["percent_covered"], + executed_lines=coverage_data[resolved_name]["executed_lines"], + unexecuted_lines=coverage_data[resolved_name]["missing_lines"], + executed_branches=coverage_data[resolved_name]["executed_branches"], + unexecuted_branches=coverage_data[resolved_name]["missing_branches"], + ) + except KeyError: + main_function_coverage = FunctionCoverage( + name=resolved_name, + coverage=0, + executed_lines=[], + unexecuted_lines=[], + executed_branches=[], + unexecuted_branches=[], + ) + + dependent_function = extract_dependent_function(function_name, code_context) + dependent_func_coverage = ( + CoverageData.grab_dependent_function_from_coverage_data( + dependent_function, coverage_data, original_cov_data + ) + if dependent_function + else None + ) + + return main_function_coverage, dependent_func_coverage + + @staticmethod + def _aggregate_coverage( + main_func_coverage: FunctionCoverage, dependent_func_coverage: Union[FunctionCoverage, None] + ) -> tuple[set[int], set[int]]: + total_executed_lines = set(main_func_coverage.executed_lines) + total_unexecuted_lines = set(main_func_coverage.unexecuted_lines) + + if dependent_func_coverage: + total_executed_lines.update(dependent_func_coverage.executed_lines) + total_unexecuted_lines.update(dependent_func_coverage.unexecuted_lines) + + return total_executed_lines, total_unexecuted_lines + + @staticmethod + def _build_graph( + main_func_coverage: FunctionCoverage, dependent_func_coverage: Union[FunctionCoverage, None] + ) -> dict[str, dict[str, Collection[object]]]: + graph = { + main_func_coverage.name: { + "executed_lines": set(main_func_coverage.executed_lines), + "unexecuted_lines": set(main_func_coverage.unexecuted_lines), + "executed_branches": main_func_coverage.executed_branches, + "unexecuted_branches": main_func_coverage.unexecuted_branches, + } + } + + if dependent_func_coverage: + graph[dependent_func_coverage.name] = { + "executed_lines": set(dependent_func_coverage.executed_lines), + "unexecuted_lines": set(dependent_func_coverage.unexecuted_lines), + "executed_branches": dependent_func_coverage.executed_branches, + "unexecuted_branches": dependent_func_coverage.unexecuted_branches, + } + + return graph + + @staticmethod + def grab_dependent_function_from_coverage_data( + dependent_function_name: str, + coverage_data: dict[str, dict[str, Any]], + original_cov_data: dict[str, dict[str, Any]], + ) -> FunctionCoverage: + """Grab the dependent function from the coverage data.""" + try: + return FunctionCoverage( + name=dependent_function_name, + coverage=coverage_data[dependent_function_name]["summary"]["percent_covered"], + executed_lines=coverage_data[dependent_function_name]["executed_lines"], + unexecuted_lines=coverage_data[dependent_function_name]["missing_lines"], + executed_branches=coverage_data[dependent_function_name]["executed_branches"], + unexecuted_branches=coverage_data[dependent_function_name]["missing_branches"], + ) + except KeyError: + msg = f"Coverage data not found for dependent function {dependent_function_name} in the coverage data" + try: + files = original_cov_data["files"] + for file in files: + functions = files[file]["functions"] + for function in functions: + if dependent_function_name in function: + return FunctionCoverage( + name=dependent_function_name, + coverage=functions[function]["summary"]["percent_covered"], + executed_lines=functions[function]["executed_lines"], + unexecuted_lines=functions[function]["missing_lines"], + executed_branches=functions[function]["executed_branches"], + unexecuted_branches=functions[function]["missing_branches"], + ) + msg = f"Coverage data not found for dependent function {dependent_function_name} in the original coverage data" + except KeyError: + raise ValueError(msg) from None + + return FunctionCoverage( + name=dependent_function_name, + coverage=0, + executed_lines=[], + unexecuted_lines=[], + executed_branches=[], + unexecuted_branches=[], + ) + def build_message(self) -> str: if self.status == CoverageStatus.NOT_FOUND: return f"No coverage data found for {self.function_name}" @@ -318,237 +555,3 @@ class FunctionCoverage: class TestingMode(enum.Enum): BEHAVIOR = "behavior" PERFORMANCE = "performance" - LINE_PROFILE = "line_profile" - - -class VerificationType(str, Enum): - FUNCTION_CALL = ( - "function_call" # Correctness verification for a test function, checks input values and output values) - ) - INIT_STATE_FTO = "init_state_fto" # Correctness verification for fto class instance attributes after init - INIT_STATE_HELPER = "init_state_helper" # Correctness verification for helper class instance attributes after init - - -class TestType(Enum): - EXISTING_UNIT_TEST = 1 - INSPIRED_REGRESSION = 2 - GENERATED_REGRESSION = 3 - REPLAY_TEST = 4 - CONCOLIC_COVERAGE_TEST = 5 - INIT_STATE_TEST = 6 - - def to_name(self) -> str: - if self is TestType.INIT_STATE_TEST: - return "" - names = { - TestType.EXISTING_UNIT_TEST: "βš™οΈ Existing Unit Tests", - TestType.INSPIRED_REGRESSION: "🎨 Inspired Regression Tests", - TestType.GENERATED_REGRESSION: "πŸŒ€ Generated Regression Tests", - TestType.REPLAY_TEST: "βͺ Replay Tests", - TestType.CONCOLIC_COVERAGE_TEST: "πŸ”Ž Concolic Coverage Tests", - } - return names[self] - - -@dataclass(frozen=True) -class InvocationId: - test_module_path: str # The fully qualified name of the test module - test_class_name: Optional[str] # The name of the class where the test is defined - test_function_name: Optional[str] # The name of the test_function. Does not include the components of the file_name - function_getting_tested: str - iteration_id: Optional[str] - - # test_module_path:TestSuiteClass.test_function_name:function_tested:iteration_id - def id(self) -> str: - class_prefix = f"{self.test_class_name}." if self.test_class_name else "" - return ( - f"{self.test_module_path}:{class_prefix}{self.test_function_name}:" - f"{self.function_getting_tested}:{self.iteration_id}" - ) - - @staticmethod - def from_str_id(string_id: str, iteration_id: Optional[str] = None) -> InvocationId: - components = string_id.split(":") - assert len(components) == 4 - second_components = components[1].split(".") - if len(second_components) == 1: - test_class_name = None - test_function_name = second_components[0] - else: - test_class_name = second_components[0] - test_function_name = second_components[1] - return InvocationId( - test_module_path=components[0], - test_class_name=test_class_name, - test_function_name=test_function_name, - function_getting_tested=components[2], - iteration_id=iteration_id if iteration_id else components[3], - ) - - -@dataclass(frozen=True) -class FunctionTestInvocation: - loop_index: int # The loop index of the function invocation, starts at 1 - id: InvocationId # The fully qualified name of the function invocation (id) - file_name: Path # The file where the test is defined - did_pass: bool # Whether the test this function invocation was part of, passed or failed - runtime: Optional[int] # Time in nanoseconds - test_framework: str # unittest or pytest - test_type: TestType - return_value: Optional[object] # The return value of the function invocation - timed_out: Optional[bool] - verification_type: Optional[str] = VerificationType.FUNCTION_CALL - stdout: Optional[str] = None - - @property - def unique_invocation_loop_id(self) -> str: - return f"{self.loop_index}:{self.id.id()}" - - -class TestResults(BaseModel): - # don't modify these directly, use the add method - # also we don't support deletion of test results elements - caution is advised - test_results: list[FunctionTestInvocation] = [] - test_result_idx: dict[str, int] = {} - - def add(self, function_test_invocation: FunctionTestInvocation) -> None: - unique_id = function_test_invocation.unique_invocation_loop_id - if unique_id in self.test_result_idx: - if DEBUG_MODE: - logger.warning(f"Test result with id {unique_id} already exists. SKIPPING") - return - self.test_result_idx[unique_id] = len(self.test_results) - self.test_results.append(function_test_invocation) - - def merge(self, other: TestResults) -> None: - original_len = len(self.test_results) - self.test_results.extend(other.test_results) - for k, v in other.test_result_idx.items(): - if k in self.test_result_idx: - msg = f"Test result with id {k} already exists." - raise ValueError(msg) - self.test_result_idx[k] = v + original_len - - def get_by_unique_invocation_loop_id(self, unique_invocation_loop_id: str) -> FunctionTestInvocation | None: - try: - return self.test_results[self.test_result_idx[unique_invocation_loop_id]] - except (IndexError, KeyError): - return None - - def get_all_ids(self) -> set[InvocationId]: - return {test_result.id for test_result in self.test_results} - - def get_all_unique_invocation_loop_ids(self) -> set[str]: - return {test_result.unique_invocation_loop_id for test_result in self.test_results} - - def number_of_loops(self) -> int: - if not self.test_results: - return 0 - return max(test_result.loop_index for test_result in self.test_results) - - def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]: - report = {} - for test_type in TestType: - report[test_type] = {"passed": 0, "failed": 0} - for test_result in self.test_results: - if test_result.loop_index == 1: - if test_result.did_pass: - report[test_result.test_type]["passed"] += 1 - else: - report[test_result.test_type]["failed"] += 1 - return report - - @staticmethod - def report_to_string(report: dict[TestType, dict[str, int]]) -> str: - return " ".join( - [ - f"{test_type.to_name()}- (Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']})" - for test_type in TestType - ] - ) - - @staticmethod - def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree: - tree = Tree(title) - for test_type in TestType: - if test_type is TestType.INIT_STATE_TEST: - continue - tree.add( - f"{test_type.to_name()} - Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']}" - ) - return tree - - def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: - for result in self.test_results: - if result.did_pass and not result.runtime: - msg = ( - f"Ignoring test case that passed but had no runtime -> {result.id}, " - f"Loop # {result.loop_index}, Test Type: {result.test_type}, " - f"Verification Type: {result.verification_type}" - ) - logger.debug(msg) - - usable_runtimes = [ - (result.id, result.runtime) for result in self.test_results if result.did_pass and result.runtime - ] - return { - usable_id: [runtime[1] for runtime in usable_runtimes if runtime[0] == usable_id] - for usable_id in {runtime[0] for runtime in usable_runtimes} - } - - def total_passed_runtime(self) -> int: - """Calculate the sum of runtimes of all test cases that passed. - - A testcase runtime is the minimum value of all looped execution runtimes. - - :return: The runtime in nanoseconds. - """ - return sum( - [min(usable_runtime_data) for _, usable_runtime_data in self.usable_runtime_data_by_test_case().items()] - ) - - def __iter__(self) -> Iterator[FunctionTestInvocation]: - return iter(self.test_results) - - def __len__(self) -> int: - return len(self.test_results) - - def __getitem__(self, index: int) -> FunctionTestInvocation: - return self.test_results[index] - - def __setitem__(self, index: int, value: FunctionTestInvocation) -> None: - self.test_results[index] = value - - def __contains__(self, value: FunctionTestInvocation) -> bool: - return value in self.test_results - - def __bool__(self) -> bool: - return bool(self.test_results) - - def __eq__(self, other: object) -> bool: - # Unordered comparison - if type(self) is not type(other): - return False - if len(self) != len(other): - return False - original_recursion_limit = sys.getrecursionlimit() - cast(TestResults, other) - for test_result in self: - other_test_result = other.get_by_unique_invocation_loop_id(test_result.unique_invocation_loop_id) - if other_test_result is None: - return False - - if original_recursion_limit < 5000: - sys.setrecursionlimit(5000) - if ( - test_result.file_name != other_test_result.file_name - or test_result.did_pass != other_test_result.did_pass - or test_result.runtime != other_test_result.runtime - or test_result.test_framework != other_test_result.test_framework - or test_result.test_type != other_test_result.test_type - or not comparator(test_result.return_value, other_test_result.return_value) - ): - sys.setrecursionlimit(original_recursion_limit) - return False - sys.setrecursionlimit(original_recursion_limit) - return True diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 953a12028..6bba42191 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -19,6 +19,7 @@ from rich.tree import Tree from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient +from codeflash.benchmarking.utils import process_benchmark_data from codeflash.cli_cmds.console import code_print, console, logger, progress_bar from codeflash.code_utils import env_utils from codeflash.code_utils.code_extractor import add_needed_imports_from_module, extract_code @@ -268,6 +269,13 @@ def optimize_function(self) -> Result[BestOptimization, str]: best_optimization.candidate.explanation, title="Best Candidate Explanation", border_style="blue" ) ) + processed_benchmark_info = None + if self.args.benchmark: + processed_benchmark_info = process_benchmark_data( + replay_performance_gain=best_optimization.replay_performance_gain, + fto_benchmark_timings=self.function_benchmark_timings, + total_benchmark_timings=self.total_benchmark_timings + ) explanation = Explanation( raw_explanation_message=best_optimization.candidate.explanation, winning_behavioral_test_results=best_optimization.winning_behavioral_test_results, @@ -276,9 +284,7 @@ def optimize_function(self) -> Result[BestOptimization, str]: best_runtime_ns=best_optimization.runtime, function_name=function_to_optimize_qualified_name, file_path=self.function_to_optimize.file_path, - replay_performance_gain=best_optimization.replay_performance_gain if self.args.benchmark else None, - fto_benchmark_timings = self.function_benchmark_timings if self.args.benchmark else None, - total_benchmark_timings = self.total_benchmark_timings if self.args.benchmark else None, + benchmark_details=processed_benchmark_info.benchmark_details if processed_benchmark_info else None ) self.log_successful_optimization(explanation, generated_tests) diff --git a/codeflash/result/create_pr.py b/codeflash/result/create_pr.py index e2d4da13c..da0c61961 100644 --- a/codeflash/result/create_pr.py +++ b/codeflash/result/create_pr.py @@ -77,6 +77,7 @@ def check_create_pr( speedup_pct=explanation.speedup_pct, winning_behavioral_test_results=explanation.winning_behavioral_test_results, winning_benchmarking_test_results=explanation.winning_benchmarking_test_results, + benchmark_details=explanation.benchmark_details ), existing_tests=existing_tests_source, generated_tests=generated_original_test_source, @@ -123,6 +124,7 @@ def check_create_pr( speedup_pct=explanation.speedup_pct, winning_behavioral_test_results=explanation.winning_behavioral_test_results, winning_benchmarking_test_results=explanation.winning_benchmarking_test_results, + benchmark_details=explanation.benchmark_details ), existing_tests=existing_tests_source, generated_tests=generated_original_test_source, diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index c44282e88..10794991a 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -5,7 +5,8 @@ from pydantic.dataclasses import dataclass from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.models.models import TestResults +from codeflash.models.models import BenchmarkDetail +from codeflash.verification.test_results import TestResults @dataclass(frozen=True, config={"arbitrary_types_allowed": True}) @@ -17,9 +18,7 @@ class Explanation: best_runtime_ns: int function_name: str file_path: Path - replay_performance_gain: Optional[float] - fto_benchmark_timings: Optional[dict[str, int]] - total_benchmark_timings: Optional[dict[str, int]] + benchmark_details: Optional[list[BenchmarkDetail]] = None @property def perf_improvement_line(self) -> str: @@ -43,29 +42,13 @@ def to_console_string(self) -> str: original_runtime_human = humanize_runtime(self.original_runtime_ns) best_runtime_human = humanize_runtime(self.best_runtime_ns) benchmark_info = "" - if self.replay_performance_gain and self.fto_benchmark_timings and self.total_benchmark_timings: - benchmark_info += "Benchmark Performance Details:\n" - for benchmark_key, og_benchmark_timing in self.fto_benchmark_timings.items(): - # benchmark key is benchmark filename :: benchmark test function :: line number - try: - benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") - except ValueError: - benchmark_info += f"Benchmark key {benchmark_key} is not in the expected format.\n" - continue - total_benchmark_timing = self.total_benchmark_timings[benchmark_key] - if total_benchmark_timing == 0: - benchmark_info += f"Benchmark timing for {benchmark_file_name}::{benchmark_test_function} was improved, but the speedup cannot be estimated.\n" - else: - # find out expected new benchmark timing, then calculate how much total benchmark was sped up. print out intermediate values - benchmark_info += f"Original timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(total_benchmark_timing)}\n" - replay_speedup = self.replay_performance_gain - expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + 1 / ( - replay_speedup + 1) * og_benchmark_timing - benchmark_info += f"Expected new timing for {benchmark_file_name}::{benchmark_test_function}: {humanize_runtime(int(expected_new_benchmark_timing))}\n" - benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing - benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 - benchmark_info += f"Benchmark speedup for {benchmark_file_name}::{benchmark_test_function}: {benchmark_speedup_percent:.2f}%\n\n" + if self.benchmark_details: + benchmark_info += "Benchmark Performance Details:\n" + for detail in self.benchmark_details: + benchmark_info += f"Original timing for {detail.benchmark_name}::{detail.test_function}: {detail.original_timing}\n" + benchmark_info += f"Expected new timing for {detail.benchmark_name}::{detail.test_function}: {detail.expected_new_timing}\n" + benchmark_info += f"Benchmark speedup for {detail.benchmark_name}::{detail.test_function}: {detail.speedup_percent:.2f}%\n\n" return ( f"Optimized {self.function_name} in {self.file_path}\n" From 4c61de9b42fff9b1583096e19102edfef3b41b2d Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 26 Mar 2025 12:08:52 -0700 Subject: [PATCH 081/122] mypy fix --- codeflash/github/PrComment.py | 11 +++-------- 1 file changed, 3 insertions(+), 8 deletions(-) diff --git a/codeflash/github/PrComment.py b/codeflash/github/PrComment.py index 4ef162cda..5b891b8a5 100644 --- a/codeflash/github/PrComment.py +++ b/codeflash/github/PrComment.py @@ -22,14 +22,14 @@ class PrComment: winning_benchmarking_test_results: TestResults benchmark_details: Optional[list[BenchmarkDetail]] = None - def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str, Optional[list[dict[str, any]]]]]: + def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str, Optional[list[BenchmarkDetail]]]]: report_table = { test_type.to_name(): result for test_type, result in self.winning_behavioral_test_results.get_test_pass_fail_report_by_type().items() if test_type.to_name() } - result = { + return { "optimization_explanation": self.optimization_explanation, "best_runtime": humanize_runtime(self.best_runtime), "original_runtime": humanize_runtime(self.original_runtime), @@ -39,14 +39,9 @@ def to_json(self) -> dict[str, Union[dict[str, dict[str, int]], int, str, Option "speedup_pct": self.speedup_pct, "loop_count": self.winning_benchmarking_test_results.number_of_loops(), "report_table": report_table, + "benchmark_details": self.benchmark_details if self.benchmark_details else None, } - # Add benchmark details if available - if self.benchmark_details: - result["benchmark_details"] = self.benchmark_details - - return result - class FileDiffContent(BaseModel): oldContent: str From eda0d46ee2cd16c883ef3c309f6f24d13877e113 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 26 Mar 2025 15:46:49 -0700 Subject: [PATCH 082/122] use dill instead of pickle --- codeflash/benchmarking/codeflash_trace.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index f708d752f..3b55aa6ba 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -1,6 +1,6 @@ import functools import os -import pickle +import dill as pickle import time from typing import Callable @@ -63,13 +63,11 @@ def wrapper(*args, **kwargs): overhead_end_time = time.thread_time_ns() overhead_time = overhead_end_time - overhead_start_time - self.function_calls_data.append( (func.__name__, class_name, func.__module__, func.__code__.co_filename, benchmark_function_name, benchmark_file_name, benchmark_line_number, execution_time, overhead_time, pickled_args, pickled_kwargs) ) - print("appended") except Exception as e: print(f"Error in codeflash_trace: {e}") From a82e9f0d49f23f57d5d7e5d0e289c5f156876411 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 28 Mar 2025 13:55:59 -0700 Subject: [PATCH 083/122] modified the benchmarking approach. codeflash_trace and codeflash_benchmark_plugins are singleton instances that will both handle writing to disk. enables flushing to disk once a limit is reached. also added various details to the tracer --- .../benchmarking/benchmark_database_utils.py | 296 ------------------ codeflash/benchmarking/codeflash_trace.py | 167 +++++++--- .../instrument_codeflash_trace.py | 85 +++-- codeflash/benchmarking/plugin/plugin.py | 262 ++++++++++++++-- .../pytest_new_process_trace_benchmarks.py | 16 +- codeflash/benchmarking/replay_test.py | 12 +- codeflash/benchmarking/trace_benchmarks.py | 5 +- codeflash/benchmarking/utils.py | 20 +- codeflash/cli_cmds/cli.py | 1 - codeflash/discovery/functions_to_optimize.py | 1 - codeflash/models/models.py | 11 + codeflash/optimization/function_optimizer.py | 19 +- codeflash/optimization/optimizer.py | 20 +- codeflash/verification/test_results.py | 93 +++--- tests/test_instrument_codeflash_trace.py | 243 +++++++++++++- tests/test_trace_benchmarks.py | 14 +- 16 files changed, 737 insertions(+), 528 deletions(-) delete mode 100644 codeflash/benchmarking/benchmark_database_utils.py diff --git a/codeflash/benchmarking/benchmark_database_utils.py b/codeflash/benchmarking/benchmark_database_utils.py deleted file mode 100644 index 1c117553c..000000000 --- a/codeflash/benchmarking/benchmark_database_utils.py +++ /dev/null @@ -1,296 +0,0 @@ -import sqlite3 -from pathlib import Path - -import pickle - - -class BenchmarkDatabaseUtils: - def __init__(self, trace_path :Path) -> None: - self.trace_path = trace_path - self.connection = None - - def setup(self) -> None: - try: - # Open connection - self.connection = sqlite3.connect(self.trace_path) - cur = self.connection.cursor() - cur.execute("PRAGMA synchronous = OFF") - cur.execute( - "CREATE TABLE IF NOT EXISTS function_calls(" - "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," - "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," - "time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" - ) - cur.execute( - "CREATE TABLE IF NOT EXISTS benchmark_timings(" - "benchmark_file_name TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER," - "time_ns INTEGER)" # Added closing parenthesis - ) - self.connection.commit() - # Don't close the connection here - except Exception as e: - print(f"Database setup error: {e}") - if self.connection: - self.connection.close() - self.connection = None - raise - - def write_function_timings(self, data: list[tuple]) -> None: - if not self.connection: - self.connection = sqlite3.connect(self.trace_path) - - try: - cur = self.connection.cursor() - # Insert data into the function_calls table - cur.executemany( - "INSERT INTO function_calls " - "(function_name, class_name, module_name, file_name, benchmark_function_name, " - "benchmark_file_name, benchmark_line_number, time_ns, overhead_time_ns, args, kwargs) " - "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", - data - ) - self.connection.commit() - except Exception as e: - print(f"Error writing to function timings database: {e}") - self.connection.rollback() - raise - - def write_benchmark_timings(self, data: list[tuple]) -> None: - if not self.connection: - self.connection = sqlite3.connect(self.trace_path) - - try: - cur = self.connection.cursor() - # Insert data into the benchmark_timings table - cur.executemany( - "INSERT INTO benchmark_timings (benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns) VALUES (?, ?, ?, ?)", - data - ) - self.connection.commit() - except Exception as e: - print(f"Error writing to benchmark timings database: {e}") - self.connection.rollback() - raise - - def print_function_timings(self, limit: int = None) -> None: - """Print the contents of a CodeflashTrace SQLite database. - - Args: - limit: Maximum number of records to print (None for all) - """ - if not self.connection: - self.connection = sqlite3.connect(self.trace_path) - try: - cur = self.connection.cursor() - - # Get the count of records - cur.execute("SELECT COUNT(*) FROM function_calls") - total_records = cur.fetchone()[0] - print(f"Found {total_records} function call records in {self.trace_path}") - - # Build the query with optional limit - query = "SELECT * FROM function_calls" - if limit: - query += f" LIMIT {limit}" - - # Execute the query - cur.execute(query) - - # Print column names - columns = [desc[0] for desc in cur.description] - print("\nColumns:", columns) - print("\n" + "=" * 80 + "\n") - - # Print each row - for i, row in enumerate(cur.fetchall()): - print(f"Record #{i + 1}:") - print(f" Function: {row[0]}") - print(f" Class: {row[1]}") - print(f" Module: {row[2]}") - print(f" File: {row[3]}") - print(f" Benchmark Function: {row[4] or 'N/A'}") - print(f" Benchmark File: {row[5] or 'N/A'}") - print(f" Benchmark Line: {row[6] or 'N/A'}") - print(f" Execution Time: {row[7]:.6f} seconds") - print(f" Overhead Time: {row[8]:.6f} seconds") - - # Unpickle and print args and kwargs - try: - args = pickle.loads(row[9]) - kwargs = pickle.loads(row[10]) - - print(f" Args: {args}") - print(f" Kwargs: {kwargs}") - except Exception as e: - print(f" Error unpickling args/kwargs: {e}") - print(f" Raw args: {row[9]}") - print(f" Raw kwargs: {row[10]}") - - print("\n" + "-" * 40 + "\n") - - except Exception as e: - print(f"Error reading database: {e}") - - def print_benchmark_timings(self, limit: int = None) -> None: - """Print the contents of a CodeflashTrace SQLite database. - Args: - limit: Maximum number of records to print (None for all) - """ - if not self.connection: - self.connection = sqlite3.connect(self.trace_path) - try: - cur = self.connection.cursor() - - # Get the count of records - cur.execute("SELECT COUNT(*) FROM benchmark_timings") - total_records = cur.fetchone()[0] - print(f"Found {total_records} benchmark timing records in {self.trace_path}") - - # Build the query with optional limit - query = "SELECT * FROM benchmark_timings" - if limit: - query += f" LIMIT {limit}" - - # Execute the query - cur.execute(query) - - # Print column names - columns = [desc[0] for desc in cur.description] - print("\nColumns:", columns) - print("\n" + "=" * 80 + "\n") - - # Print each row - for i, row in enumerate(cur.fetchall()): - print(f"Record #{i + 1}:") - print(f" Benchmark File: {row[0] or 'N/A'}") - print(f" Benchmark Function: {row[1] or 'N/A'}") - print(f" Benchmark Line: {row[2] or 'N/A'}") - print(f" Execution Time: {row[3] / 1e9:.6f} seconds") # Convert nanoseconds to seconds - print("\n" + "-" * 40 + "\n") - - except Exception as e: - print(f"Error reading benchmark timings database: {e}") - - - def close(self) -> None: - if self.connection: - self.connection.close() - self.connection = None - - - @staticmethod - def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[str, int]]: - """Process the trace file and extract timing data for all functions. - - Args: - trace_path: Path to the trace file - - Returns: - A nested dictionary where: - - Outer keys are module_name.qualified_name (module.class.function) - - Inner keys are benchmark filename :: benchmark test function :: line number - - Values are function timing in milliseconds - - """ - # Initialize the result dictionary - result = {} - - # Connect to the SQLite database - connection = sqlite3.connect(trace_path) - cursor = connection.cursor() - - try: - # Query the function_calls table for all function calls - cursor.execute( - "SELECT module_name, class_name, function_name, " - "benchmark_file_name, benchmark_function_name, benchmark_line_number, time_ns " - "FROM function_calls" - ) - - # Process each row - for row in cursor.fetchall(): - module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row - - # Create the function key (module_name.class_name.function_name) - if class_name: - qualified_name = f"{module_name}.{class_name}.{function_name}" - else: - qualified_name = f"{module_name}.{function_name}" - - # Create the benchmark key (file::function::line) - benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - - # Initialize the inner dictionary if needed - if qualified_name not in result: - result[qualified_name] = {} - - # If multiple calls to the same function in the same benchmark, - # add the times together - if benchmark_key in result[qualified_name]: - result[qualified_name][benchmark_key] += time_ns - else: - result[qualified_name][benchmark_key] = time_ns - - finally: - # Close the connection - connection.close() - - return result - - @staticmethod - def get_benchmark_timings(trace_path: Path) -> dict[str, int]: - """Extract total benchmark timings from trace files. - - Args: - trace_path: Path to the trace file - - Returns: - A dictionary mapping where: - - Keys are benchmark filename :: benchmark test function :: line number - - Values are total benchmark timing in milliseconds (with overhead subtracted) - - """ - # Initialize the result dictionary - result = {} - overhead_by_benchmark = {} - - # Connect to the SQLite database - connection = sqlite3.connect(trace_path) - cursor = connection.cursor() - - try: - # Query the function_calls table to get total overhead for each benchmark - cursor.execute( - "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " - "FROM function_calls " - "GROUP BY benchmark_file_name, 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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - 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_name, benchmark_function_name, benchmark_line_number, time_ns " - "FROM benchmark_timings" - ) - - # Process each row and subtract overhead - for row in cursor.fetchall(): - benchmark_file, benchmark_func, benchmark_line, time_ns = row - - # Create the benchmark key (file::function::line) - benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - - # Subtract overhead from total time - overhead = overhead_by_benchmark.get(benchmark_key, 0) - result[benchmark_key] = time_ns - overhead - - finally: - # Close the connection - connection.close() - - return result diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 3b55aa6ba..2ae57307b 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -1,23 +1,90 @@ import functools import os -import dill as pickle -import time -from typing import Callable - +import sqlite3 +import sys +import pickle +import dill +import time +from typing import Callable, Optional class CodeflashTrace: - """A class that provides both a decorator for tracing function calls - and a context manager for managing the tracing data lifecycle. - """ + """Decorator class that traces and profiles function execution.""" def __init__(self) -> None: self.function_calls_data = [] + self.function_call_count = 0 + self.pickle_count_limit = 1000 + self._connection = None + self._trace_path = None + + def setup(self, trace_path: str) -> None: + """Set up the database connection for direct writing. + + Args: + trace_path: Path to the trace database file - def __exit__(self, exc_type, exc_val, exc_tb) -> None: - # Cleanup is optional here - pass + """ + try: + self._trace_path = trace_path + self._connection = sqlite3.connect(self._trace_path) + cur = self._connection.cursor() + cur.execute("PRAGMA synchronous = OFF") + cur.execute( + "CREATE TABLE IF NOT EXISTS benchmark_function_timings(" + "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," + "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," + "function_time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" + ) + self._connection.commit() + except Exception as e: + print(f"Database setup error: {e}") + if self._connection: + self._connection.close() + self._connection = None + raise + + def write_function_timings(self) -> None: + """Write function call data directly to the database. + + Args: + data: List of function call data tuples to write + """ + if not self.function_calls_data: + return # No data to write + + if self._connection is None and self._trace_path is not None: + self._connection = sqlite3.connect(self._trace_path) + + try: + cur = self._connection.cursor() + # Insert data into the benchmark_function_timings table + cur.executemany( + "INSERT INTO benchmark_function_timings" + "(function_name, class_name, module_name, file_name, benchmark_function_name, " + "benchmark_file_name, benchmark_line_number, function_time_ns, overhead_time_ns, args, kwargs) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + self.function_calls_data + ) + self._connection.commit() + self.function_calls_data = [] + except Exception as e: + print(f"Error writing to function timings database: {e}") + if self._connection: + self._connection.rollback() + raise + + def open(self) -> None: + """Open the database connection.""" + if self._connection is None: + self._connection = sqlite3.connect(self._trace_path) + + def close(self) -> None: + """Close the database connection.""" + if self._connection: + self._connection.close() + self._connection = None def __call__(self, func: Callable) -> Callable: """Use as a decorator to trace function execution. @@ -38,39 +105,55 @@ def wrapper(*args, **kwargs): # Calculate execution time execution_time = end_time - start_time - # Measure overhead - overhead_start_time = time.thread_time_ns() - - try: - # Check if currently in pytest benchmark fixture - if os.environ.get("CODEFLASH_BENCHMARKING", "False") == "False": - return result - - # 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", "") - benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "") - # Get class name - class_name = "" - qualname = func.__qualname__ - if "." in qualname: - class_name = qualname.split(".")[0] - # Calculate overhead time - overhead_end_time = time.thread_time_ns() - overhead_time = overhead_end_time - overhead_start_time - - self.function_calls_data.append( - (func.__name__, class_name, func.__module__, func.__code__.co_filename, - benchmark_function_name, benchmark_file_name, benchmark_line_number, execution_time, - overhead_time, pickled_args, pickled_kwargs) - ) - except Exception as e: - print(f"Error in codeflash_trace: {e}") + 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": + return result + + # 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", "") + benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "") + # Get class name + class_name = "" + qualname = func.__qualname__ + if "." in qualname: + class_name = qualname.split(".")[0] + + if self.function_call_count <= self.pickle_count_limit: + 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) + 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 + + if len(self.function_calls_data) > 1000: + self.write_function_timings() + # Calculate overhead time + 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_file_name, benchmark_line_number, execution_time, + overhead_time, pickled_args, pickled_kwargs) + ) return result return wrapper diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py index 017cecaca..06e93daf8 100644 --- a/codeflash/benchmarking/instrument_codeflash_trace.py +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -1,3 +1,5 @@ +from pathlib import Path + import isort import libcst as cst @@ -5,40 +7,35 @@ class AddDecoratorTransformer(cst.CSTTransformer): - def __init__(self, function_name, class_name=None): + def __init__(self, target_functions: set[tuple[str, str]]) -> None: super().__init__() - self.function_name = function_name - self.class_name = class_name - self.in_target_class = (class_name is None) # If no class name, always "in target class" + self.target_functions = target_functions self.added_codeflash_trace = False + self.class_name = "" + self.decorator = cst.Decorator( + decorator=cst.Name(value="codeflash_trace") + ) def leave_ClassDef(self, original_node, updated_node): - if self.class_name and original_node.name.value == self.class_name: - self.in_target_class = False + self.class_name = "" return updated_node def visit_ClassDef(self, node): - if self.class_name and node.name.value == self.class_name: - self.in_target_class = True - return True + if self.class_name: # Don't go into nested class + return False + self.class_name = node.name.value def leave_FunctionDef(self, original_node, updated_node): - if not self.in_target_class or original_node.name.value != self.function_name: + if (self.class_name, original_node.name.value) in self.target_functions: + # Add the new decorator after any existing decorators, so it gets executed first + updated_decorators = list(updated_node.decorators) + [self.decorator] + self.added_codeflash_trace = True + return updated_node.with_changes( + decorators=updated_decorators + ) + else: return updated_node - # Create the codeflash_trace decorator - decorator = cst.Decorator( - decorator=cst.Name(value="codeflash_trace") - ) - - # Add the new decorator after any existing decorators - updated_decorators = list(updated_node.decorators) + [decorator] - self.added_codeflash_trace = True - # Return the updated node with the new decorator - return updated_node.with_changes( - decorators=updated_decorators - ) - def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: # Create import statement for codeflash_trace if not self.added_codeflash_trace: @@ -62,12 +59,12 @@ def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> c ] ) - # Insert at the beginning of the file + # Insert at the beginning of the file. We'll use isort later to sort the imports. new_body = [import_stmt, *list(updated_node.body)] return updated_node.with_changes(body=new_body) -def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToOptimize) -> str: +def add_codeflash_decorator_to_code(code: str, functions_to_optimize: list[FunctionToOptimize]) -> str: """Add codeflash_trace to a function. Args: @@ -76,15 +73,17 @@ def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToO Returns: The modified source code as a string + """ - # Extract class name if present - class_name = None - if len(function_to_optimize.parents) == 1 and function_to_optimize.parents[0].type == "ClassDef": - class_name = function_to_optimize.parents[0].name + target_functions = set() + for function_to_optimize in functions_to_optimize: + class_name = "" + if len(function_to_optimize.parents) == 1 and function_to_optimize.parents[0].type == "ClassDef": + class_name = function_to_optimize.parents[0].name + target_functions.add((class_name, function_to_optimize.function_name)) transformer = AddDecoratorTransformer( - function_name=function_to_optimize.function_name, - class_name=class_name + target_functions = target_functions, ) module = cst.parse_module(code) @@ -93,17 +92,17 @@ def add_codeflash_decorator_to_code(code: str, function_to_optimize: FunctionToO def instrument_codeflash_trace_decorator( - function_to_optimize: FunctionToOptimize + file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]] ) -> None: - """Instrument __init__ function with codeflash_trace decorator if it's in a class.""" - # Instrument fto class - original_code = function_to_optimize.file_path.read_text(encoding="utf-8") - new_code = add_codeflash_decorator_to_code( - original_code, - function_to_optimize - ) - # Modify the code - modified_code = isort.code(code=new_code, float_to_top=True) + """Instrument codeflash_trace decorator to functions to optimize.""" + for file_path, functions_to_optimize in file_to_funcs_to_optimize.items(): + original_code = file_path.read_text(encoding="utf-8") + new_code = add_codeflash_decorator_to_code( + original_code, + functions_to_optimize + ) + # Modify the code + modified_code = isort.code(code=new_code, float_to_top=True) - # Write the modified code back to the file - function_to_optimize.file_path.write_text(modified_code, encoding="utf-8") + # Write the modified code back to the file + file_path.write_text(modified_code, encoding="utf-8") diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index ee7504ec4..9d7da6ef2 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -1,33 +1,195 @@ +from __future__ import annotations +import os +import sqlite3 import sys - -import pytest import time -import os +from pathlib import Path +import pytest +from codeflash.benchmarking.codeflash_trace import codeflash_trace +from codeflash.models.models import BenchmarkKey + + class CodeFlashBenchmarkPlugin: - benchmark_timings = [] + def __init__(self) -> None: + self._trace_path = None + self._connection = None + self.benchmark_timings = [] - class Benchmark: - def __init__(self, request): - self.request = request + def setup(self, trace_path:str) -> None: + try: + # Open connection + self._trace_path = trace_path + self._connection = sqlite3.connect(self._trace_path) + cur = self._connection.cursor() + cur.execute("PRAGMA synchronous = OFF") + cur.execute( + "CREATE TABLE IF NOT EXISTS benchmark_timings(" + "benchmark_file_name TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER," + "benchmark_time_ns INTEGER)" + ) + self._connection.commit() + self.close() # Reopen only at the end of pytest session + except Exception as e: + print(f"Database setup error: {e}") + if self._connection: + self._connection.close() + self._connection = None + raise - def __call__(self, func, *args, **kwargs): - benchmark_file_name = self.request.node.fspath.basename - benchmark_function_name = self.request.node.name - line_number = str(sys._getframe(1).f_lineno) # 1 frame up in the call stack + def write_benchmark_timings(self) -> None: + if not self.benchmark_timings: + return # No data to write - os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = benchmark_function_name - os.environ["CODEFLASH_BENCHMARK_FILE_NAME"] = benchmark_file_name - os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = line_number - os.environ["CODEFLASH_BENCHMARKING"] = "True" + if self._connection is None: + self._connection = sqlite3.connect(self._trace_path) - start = time.perf_counter_ns() - result = func(*args, **kwargs) - end = time.perf_counter_ns() + try: + cur = self._connection.cursor() + # Insert data into the benchmark_timings table + cur.executemany( + "INSERT INTO benchmark_timings (benchmark_file_name, benchmark_function_name, benchmark_line_number, benchmark_time_ns) VALUES (?, ?, ?, ?)", + self.benchmark_timings + ) + self._connection.commit() + self.benchmark_timings = [] # Clear the benchmark timings list + except Exception as e: + print(f"Error writing to benchmark timings database: {e}") + self._connection.rollback() + raise + def close(self) -> None: + if self._connection: + self._connection.close() + self._connection = None + + @staticmethod + def get_function_benchmark_timings(trace_path: Path) -> dict[str, dict[BenchmarkKey, int]]: + """Process the trace file and extract timing data for all functions. + + Args: + trace_path: Path to the trace file + + Returns: + A nested dictionary where: + - Outer keys are module_name.qualified_name (module.class.function) + - Inner keys are of type BenchmarkKey + - Values are function timing in milliseconds + + """ + # Initialize the result dictionary + result = {} + + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + + try: + # Query the function_calls table for all function calls + cursor.execute( + "SELECT module_name, class_name, function_name, " + "benchmark_file_name, benchmark_function_name, benchmark_line_number, function_time_ns " + "FROM benchmark_function_timings" + ) + + # Process each row + for row in cursor.fetchall(): + module_name, class_name, function_name, benchmark_file, benchmark_func, benchmark_line, time_ns = row + + # Create the function key (module_name.class_name.function_name) + if class_name: + qualified_name = f"{module_name}.{class_name}.{function_name}" + else: + qualified_name = f"{module_name}.{function_name}" + + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + # Initialize the inner dictionary if needed + if qualified_name not in result: + result[qualified_name] = {} + + # If multiple calls to the same function in the same benchmark, + # add the times together + if benchmark_key in result[qualified_name]: + result[qualified_name][benchmark_key] += time_ns + else: + result[qualified_name][benchmark_key] = time_ns + + finally: + # Close the connection + connection.close() + + return result + + @staticmethod + def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]: + """Extract total benchmark timings from trace files. + + Args: + trace_path: Path to the trace file + + Returns: + A dictionary mapping where: + - Keys are of type BenchmarkKey + - Values are total benchmark timing in milliseconds (with overhead subtracted) + + """ + # Initialize the result dictionary + result = {} + overhead_by_benchmark = {} + + # Connect to the SQLite database + connection = sqlite3.connect(trace_path) + cursor = connection.cursor() + + try: + # Query the benchmark_function_timings table to get total overhead for each benchmark + cursor.execute( + "SELECT benchmark_file_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " + "FROM benchmark_function_timings " + "GROUP BY benchmark_file_name, 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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + 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_name, benchmark_function_name, benchmark_line_number, benchmark_time_ns " + "FROM benchmark_timings" + ) + + # Process each row and subtract overhead + for row in cursor.fetchall(): + benchmark_file, benchmark_func, benchmark_line, time_ns = row + + # Create the benchmark key (file::function::line) + benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" + benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + # Subtract overhead from total time + overhead = overhead_by_benchmark.get(benchmark_key, 0) + result[benchmark_key] = time_ns - overhead + + finally: + # Close the connection + connection.close() + + return result + + # Pytest hooks + @pytest.hookimpl + def pytest_sessionfinish(self, session, exitstatus): + """Execute after whole test run is completed.""" + # Write any remaining benchmark timings to the database + codeflash_trace.close() + if self.benchmark_timings: + self.write_benchmark_timings() + # Close the database connection + self.close() - os.environ["CODEFLASH_BENCHMARKING"] = "False" - CodeFlashBenchmarkPlugin.benchmark_timings.append( - (benchmark_file_name, benchmark_function_name, line_number, end - start)) - return result @staticmethod def pytest_addoption(parser): parser.addoption( @@ -39,11 +201,13 @@ def pytest_addoption(parser): @staticmethod def pytest_plugin_registered(plugin, manager): + # Not necessary since run with -p no:benchmark, but just in case if hasattr(plugin, "name") and plugin.name == "pytest-benchmark": manager.unregister(plugin) @staticmethod def pytest_collection_modifyitems(config, items): + # Skip tests that don't have the benchmark fixture if not config.getoption("--codeflash-trace"): return @@ -53,10 +217,62 @@ def pytest_collection_modifyitems(config, items): continue 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 + + """ + benchmark_file_name = self.request.node.fspath.basename + 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_NAME"] = benchmark_file_name + os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(line_number) + os.environ["CODEFLASH_BENCHMARKING"] = "True" + + # Run the function + start = time.perf_counter_ns() + result = func(*args, **kwargs) + end = time.perf_counter_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 + codeflash_trace.function_call_count = 0 + # Add to the benchmark timings buffer + codeflash_benchmark_plugin.benchmark_timings.append( + (benchmark_file_name, benchmark_function_name, line_number, end - start)) + + return result + @staticmethod @pytest.fixture def benchmark(request): if not request.config.getoption("--codeflash-trace"): return None - return CodeFlashBenchmarkPlugin.Benchmark(request) \ No newline at end of file + return CodeFlashBenchmarkPlugin.Benchmark(request) + +codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin() \ No newline at end of file diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 6d4c85f41..7b6bd747a 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -1,11 +1,8 @@ import sys from pathlib import Path -from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils -from codeflash.verification.verification_utils import get_test_file_path -from plugin.plugin import CodeFlashBenchmarkPlugin from codeflash.benchmarking.codeflash_trace import codeflash_trace -from codeflash.code_utils.code_utils import get_run_tmp_file +from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin benchmarks_root = sys.argv[1] tests_root = sys.argv[2] @@ -16,16 +13,11 @@ import pytest try: - db = BenchmarkDatabaseUtils(trace_path=Path(trace_file)) - db.setup() + codeflash_benchmark_plugin.setup(trace_file) + codeflash_trace.setup(trace_file) exitcode = pytest.main( - [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[CodeFlashBenchmarkPlugin()] + [benchmarks_root, "--codeflash-trace", "-p", "no:benchmark", "-s", "-o", "addopts="], plugins=[codeflash_benchmark_plugin] ) # Errors will be printed to stdout, not stderr - db.write_function_timings(codeflash_trace.function_calls_data) - db.write_benchmark_timings(CodeFlashBenchmarkPlugin.benchmark_timings) - # db.print_function_timings() - # db.print_benchmark_timings() - db.close() except Exception as e: print(f"Failed to collect tests: {e!s}", file=sys.stderr) diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index a1d5b370a..670d6e4bd 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -13,7 +13,7 @@ from pathlib import Path def get_next_arg_and_return( - trace_file: str, function_name: str, file_name: str, class_name: str | None = None, num_to_get: int = 25 + trace_file: str, function_name: str, file_name: str, class_name: str | None = None, num_to_get: int = 256 ) -> Generator[Any]: db = sqlite3.connect(trace_file) cur = db.cursor() @@ -21,12 +21,12 @@ def get_next_arg_and_return( if class_name is not None: cursor = cur.execute( - "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? AND class_name = ? ORDER BY time_ns ASC LIMIT ?", + "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_name = ? AND class_name = ? LIMIT ?", (function_name, file_name, class_name, limit), ) else: cursor = cur.execute( - "SELECT * FROM function_calls WHERE function_name = ? AND file_name = ? AND class_name = '' ORDER BY time_ns ASC LIMIT ?", + "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_name = ? AND class_name = '' LIMIT ?", (function_name, file_name, limit), ) @@ -42,7 +42,7 @@ def create_trace_replay_test_code( trace_file: str, functions_data: list[dict[str, Any]], test_framework: str = "pytest", - max_run_count=100 + max_run_count=256 ) -> str: """Create a replay test for functions based on trace data. @@ -217,7 +217,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework # Get distinct benchmark names cursor.execute( - "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM function_calls" + "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM benchmark_function_timings" ) benchmarks = cursor.fetchall() @@ -226,7 +226,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework benchmark_function_name, benchmark_file_name = benchmark # Get functions associated with this benchmark cursor.execute( - "SELECT DISTINCT function_name, class_name, module_name, file_name, benchmark_line_number FROM function_calls " + "SELECT DISTINCT function_name, class_name, module_name, file_name, benchmark_line_number FROM benchmark_function_timings " "WHERE benchmark_function_name = ? AND benchmark_file_name = ?", (benchmark_function_name, benchmark_file_name) ) diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 79395db79..8882078d9 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -9,7 +9,7 @@ from pathlib import Path import subprocess -def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, trace_file: Path) -> None: +def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, trace_file: Path, timeout:int = 300) -> None: result = subprocess.run( [ SAFE_SYS_EXECUTABLE, @@ -23,6 +23,7 @@ def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root capture_output=True, text=True, env={"PYTHONPATH": str(project_root)}, + timeout=timeout, ) if result.returncode != 0: if "ERROR collecting" in result.stdout: @@ -38,5 +39,5 @@ def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root else: error_section = result.stdout logger.warning( - f"Error collecting benchmarks - Pytest Exit code: {result.returncode}={ExitCode(result.returncode).name}\n {error_section}" + f"Error collecting benchmarks - Pytest Exit code: {result.returncode}, {error_section}" ) \ No newline at end of file diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 38c31b55b..5f14f141f 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -6,29 +6,30 @@ from codeflash.cli_cmds.console import logger from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.models.models import ProcessedBenchmarkInfo, BenchmarkDetail +from codeflash.models.models import ProcessedBenchmarkInfo, BenchmarkDetail, BenchmarkKey +from codeflash.result.critic import performance_gain -def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[str, int]], - total_benchmark_timings: dict[str, int]) -> dict[str, list[tuple[str, float, float, float]]]: +def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[BenchmarkKey, int]], + total_benchmark_timings: dict[BenchmarkKey, int]) -> dict[str, list[tuple[str, float, float, float]]]: function_to_result = {} # Process each function's benchmark data for func_path, test_times in function_benchmark_timings.items(): # Sort by percentage (highest first) sorted_tests = [] - for test_name, func_time in test_times.items(): - total_time = total_benchmark_timings.get(test_name, 0) + for benchmark_key, func_time in test_times.items(): + total_time = total_benchmark_timings.get(benchmark_key, 0) if func_time > total_time: - logger.debug(f"Skipping test {test_name} due to func_time {func_time} > total_time {total_time}") + logger.debug(f"Skipping test {benchmark_key} due to func_time {func_time} > total_time {total_time}") # If the function time is greater than total time, likely to have multithreading / multiprocessing issues. # Do not try to project the optimization impact for this function. - sorted_tests.append((test_name, 0.0, 0.0, 0.0)) + sorted_tests.append((str(benchmark_key), 0.0, 0.0, 0.0)) if total_time > 0: percentage = (func_time / total_time) * 100 # Convert nanoseconds to milliseconds func_time_ms = func_time / 1_000_000 total_time_ms = total_time / 1_000_000 - sorted_tests.append((test_name, total_time_ms, func_time_ms, percentage)) + sorted_tests.append((str(benchmark_key), total_time_ms, func_time_ms, percentage)) sorted_tests.sort(key=lambda x: x[3], reverse=True) function_to_result[func_path] = sorted_tests return function_to_result @@ -107,8 +108,7 @@ def process_benchmark_data( ) * og_benchmark_timing # Calculate speedup - benchmark_speedup_ratio = total_benchmark_timing / expected_new_benchmark_timing - benchmark_speedup_percent = (benchmark_speedup_ratio - 1) * 100 + benchmark_speedup_percent = performance_gain(original_runtime_ns=total_benchmark_timing, optimized_runtime_ns=int(expected_new_benchmark_timing)) * 100 benchmark_details.append( BenchmarkDetail( diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 96bb0cef3..d1e786703 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -121,7 +121,6 @@ def process_pyproject_config(args: Namespace) -> Namespace: "disable_telemetry", "disable_imports_sorting", "git_remote", - "benchmarks_root" ] for key in supported_keys: if key in pyproject_config and ( diff --git a/codeflash/discovery/functions_to_optimize.py b/codeflash/discovery/functions_to_optimize.py index cd0bfc50a..6e4f744d7 100644 --- a/codeflash/discovery/functions_to_optimize.py +++ b/codeflash/discovery/functions_to_optimize.py @@ -363,7 +363,6 @@ def visit_ClassDef(self, node: ast.ClassDef) -> None: for decorator in body_node.decorator_list ): self.is_staticmethod = True - print(f"static method found: {self.function_name}") return elif self.line_no: # If we have line number info, check if class has a static method with the same line number diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 52d1e4285..e046cf910 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -78,6 +78,15 @@ class BestOptimization(BaseModel): winning_benchmarking_test_results: TestResults winning_replay_benchmarking_test_results : Optional[TestResults] = None +@dataclass(frozen=True) +class BenchmarkKey: + file_name: str + function_name: str + line_number: int + + def __str__(self) -> str: + return f"{self.file_name}::{self.function_name}::{self.line_number}" + @dataclass class BenchmarkDetail: benchmark_name: str @@ -156,6 +165,7 @@ class OptimizedCandidateResult(BaseModel): best_test_runtime: int behavior_test_results: TestResults benchmarking_test_results: TestResults + replay_benchmarking_test_results: Optional[TestResults] = None optimization_candidate_index: int total_candidate_timing: int @@ -260,6 +270,7 @@ class FunctionParent: class OriginalCodeBaseline(BaseModel): behavioral_test_results: TestResults benchmarking_test_results: TestResults + replay_benchmarking_test_results: Optional[TestResults] = None runtime: int coverage_results: Optional[CoverageData] diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 6bba42191..464027778 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -145,8 +145,6 @@ def optimize_function(self) -> Result[BestOptimization, str]: return Failure("Codeflash does not support async functions in the code to optimize.") code_print(code_context.read_writable_code) - logger.info("Read only code") - code_print(code_context.read_only_context_code) generated_test_paths = [ get_test_file_path( self.test_cfg.tests_root, self.function_to_optimize.function_name, test_index, test_type="unit" @@ -435,8 +433,8 @@ def determine_best_candidate( tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") if self.args.benchmark: - original_code_replay_runtime = original_code_baseline.benchmarking_test_results.total_replay_test_runtime() - candidate_replay_runtime = candidate_result.benchmarking_test_results.total_replay_test_runtime() + original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results.total_passed_runtime() + candidate_replay_runtime = candidate_result.replay_benchmarking_test_results.total_passed_runtime() replay_perf_gain = performance_gain( original_runtime_ns=original_code_replay_runtime, optimized_runtime_ns=candidate_replay_runtime, @@ -951,12 +949,14 @@ def establish_original_code_baseline( logger.debug(f"Total original code runtime (ns): {total_timing}") if self.args.benchmark: - logger.info(f"Total replay test runtime: {humanize_runtime(benchmarking_results.total_replay_test_runtime())}") + replay_benchmarking_test_results = benchmarking_results.filter(TestType.REPLAY_TEST) + logger.info(f"Total replay test runtime: {humanize_runtime(replay_benchmarking_test_results.total_passed_runtime())}") return Success( ( OriginalCodeBaseline( behavioral_test_results=behavioral_results, benchmarking_test_results=benchmarking_results, + replay_benchmarking_test_results = replay_benchmarking_test_results if self.args.benchmark else None, runtime=total_timing, coverage_results=coverage_results, ), @@ -1071,13 +1071,9 @@ def run_optimized_candidate( logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") if self.args.benchmark: - total_candidate_replay_timing = ( - candidate_benchmarking_results.total_replay_test_runtime() - if candidate_benchmarking_results - else 0 - ) + candidate_replay_benchmarking_results = candidate_benchmarking_results.filter(TestType.REPLAY_TEST) logger.debug( - f"Total optimized code {optimization_candidate_index} replay benchmark runtime (ns): {total_candidate_replay_timing}" + f"Total optimized code {optimization_candidate_index} replay benchmark runtime (ns): {candidate_replay_benchmarking_results.total_passed_runtime()}" ) return Success( OptimizedCandidateResult( @@ -1085,6 +1081,7 @@ def run_optimized_candidate( best_test_runtime=total_candidate_timing, behavior_test_results=candidate_behavior_results, benchmarking_test_results=candidate_benchmarking_results, + replay_benchmarking_test_results = candidate_replay_benchmarking_results if self.args.benchmark else None, optimization_candidate_index=optimization_candidate_index, total_candidate_timing=total_candidate_timing, ) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 08db413b2..4d17a5255 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient -from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils +from codeflash.benchmarking.plugin.plugin import CodeFlashBenchmarkPlugin from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table @@ -20,7 +20,7 @@ from codeflash.discovery.discover_unit_tests import discover_unit_tests from codeflash.discovery.functions_to_optimize import get_functions_to_optimize from codeflash.either import is_successful -from codeflash.models.models import TestFiles, ValidCode +from codeflash.models.models import ValidCode, BenchmarkKey from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.telemetry.posthog_cf import ph from codeflash.verification.test_results import TestType @@ -96,8 +96,8 @@ def run(self) -> None: project_root=self.args.project_root, module_root=self.args.module_root, ) - function_benchmark_timings = None - total_benchmark_timings = None + function_benchmark_timings: dict[str, dict[BenchmarkKey, int]] = {} + total_benchmark_timings: dict[BenchmarkKey, int] = {} if self.args.benchmark: with progress_bar( f"Running benchmarks in {self.args.benchmarks_root}", @@ -109,9 +109,7 @@ def run(self) -> None: 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: - instrument_codeflash_trace_decorator(fto) + instrument_codeflash_trace_decorator(file_to_funcs_to_optimize) trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" replay_tests_dir = Path(self.args.tests_root) / "codeflash_replay_tests" trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Run all tests that use pytest-benchmark @@ -119,8 +117,8 @@ def run(self) -> None: if replay_count == 0: logger.info(f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization") else: - function_benchmark_timings = BenchmarkDatabaseUtils.get_function_benchmark_timings(trace_file) - total_benchmark_timings = BenchmarkDatabaseUtils.get_benchmark_timings(trace_file) + function_benchmark_timings = CodeFlashBenchmarkPlugin.get_function_benchmark_timings(trace_file) + total_benchmark_timings = CodeFlashBenchmarkPlugin.get_benchmark_timings(trace_file) function_to_results = validate_and_format_benchmark_table(function_benchmark_timings, total_benchmark_timings) print_benchmark_table(function_to_results) logger.info("Finished tracing existing benchmarks") @@ -191,6 +189,7 @@ def run(self) -> None: validated_original_code[analysis.file_path] = ValidCode( source_code=callee_original_code, normalized_code=normalized_callee_original_code ) + if has_syntax_error: continue @@ -200,7 +199,7 @@ def run(self) -> None: f"Optimizing function {function_iterator_count} of {num_optimizable_functions}: " f"{function_to_optimize.qualified_name}" ) - + console.rule() if not ( function_to_optimize_ast := get_first_top_level_function_or_method_ast( function_to_optimize.function_name, function_to_optimize.parents, original_module_ast @@ -256,7 +255,6 @@ def run(self) -> None: get_run_tmp_file.tmpdir.cleanup() - def run_with_args(args: Namespace) -> None: optimizer = Optimizer(args) optimizer.run() diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py index 916f6da11..6f4f397ab 100644 --- a/codeflash/verification/test_results.py +++ b/codeflash/verification/test_results.py @@ -1,10 +1,10 @@ from __future__ import annotations import sys -from collections.abc import Iterator +from collections import defaultdict from enum import Enum from pathlib import Path -from typing import Optional, cast +from typing import TYPE_CHECKING, Optional, cast from pydantic import BaseModel from pydantic.dataclasses import dataclass @@ -13,6 +13,9 @@ from codeflash.cli_cmds.console import DEBUG_MODE, logger from codeflash.verification.comparator import comparator +if TYPE_CHECKING: + from collections.abc import Iterator + class VerificationType(str, Enum): FUNCTION_CALL = ( @@ -31,7 +34,7 @@ class TestType(Enum): INIT_STATE_TEST = 6 def to_name(self) -> str: - if self == TestType.INIT_STATE_TEST: + if self is TestType.INIT_STATE_TEST: return "" names = { TestType.EXISTING_UNIT_TEST: "βš™οΈ Existing Unit Tests", @@ -53,7 +56,11 @@ class InvocationId: # test_module_path:TestSuiteClass.test_function_name:function_tested:iteration_id def id(self) -> str: - return f"{self.test_module_path}:{(self.test_class_name + '.' if self.test_class_name else '')}{self.test_function_name}:{self.function_getting_tested}:{self.iteration_id}" + class_prefix = f"{self.test_class_name}." if self.test_class_name else "" + return ( + f"{self.test_module_path}:{class_prefix}{self.test_function_name}:" + f"{self.function_getting_tested}:{self.iteration_id}" + ) @staticmethod def from_str_id(string_id: str, iteration_id: Optional[str] = None) -> InvocationId: @@ -66,7 +73,6 @@ def from_str_id(string_id: str, iteration_id: Optional[str] = None) -> Invocatio else: test_class_name = second_components[0] test_function_name = second_components[1] - # logger.debug(f"Invocation id info: test_module_path: {components[0]}, test_class_name: {test_class_name}, test_function_name: {test_function_name}, function_getting_tested: {components[2]}, iteration_id: {iteration_id if iteration_id else components[3]}") return InvocationId( test_module_path=components[0], test_class_name=test_class_name, @@ -88,6 +94,7 @@ class FunctionTestInvocation: return_value: Optional[object] # The return value of the function invocation timed_out: Optional[bool] verification_type: Optional[str] = VerificationType.FUNCTION_CALL + stdout: Optional[str] = None @property def unique_invocation_loop_id(self) -> str: @@ -118,6 +125,15 @@ def merge(self, other: TestResults) -> None: raise ValueError(msg) self.test_result_idx[k] = v + original_len + def filter(self, test_type: TestType) -> TestResults: + filtered_test_results = [] + filtered_test_results_idx = {} + for test_result in self.test_results: + if test_result.test_type == test_type: + filtered_test_results_idx[test_result.unique_invocation_loop_id] = len(filtered_test_results) + filtered_test_results.append(test_result) + return TestResults(test_results=filtered_test_results, test_result_idx=filtered_test_results_idx) + def get_by_unique_invocation_loop_id(self, unique_invocation_loop_id: str) -> FunctionTestInvocation | None: try: return self.test_results[self.test_result_idx[unique_invocation_loop_id]] @@ -160,67 +176,40 @@ def report_to_string(report: dict[TestType, dict[str, int]]) -> str: def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree: tree = Tree(title) for test_type in TestType: + if test_type is TestType.INIT_STATE_TEST: + continue tree.add( f"{test_type.to_name()} - Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']}" ) return tree def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: + usable_runtime_by_id = defaultdict(list) for result in self.test_results: - if result.did_pass and not result.runtime: - pass - # logger.debug( - # f"Ignoring test case that passed but had no runtime -> {result.id}, Loop # {result.loop_index}, Test Type: {result.test_type}, Verification Type: {result.verification_type}" - # ) - usable_runtimes = [ - (result.id, result.runtime) for result in self.test_results if result.did_pass and result.runtime - ] - return { - usable_id: [runtime[1] for runtime in usable_runtimes if runtime[0] == usable_id] - for usable_id in {runtime[0] for runtime in usable_runtimes} - } - - def total_passed_runtime(self) -> int: - """Calculate the sum of runtimes of all test cases that passed, where a testcase runtime - is the minimum value of all looped execution runtimes. + if result.did_pass: + if not result.runtime: + msg = ( + f"Ignoring test case that passed but had no runtime -> {result.id}, " + f"Loop # {result.loop_index}, Test Type: {result.test_type}, " + f"Verification Type: {result.verification_type}" + ) + logger.debug(msg) + else: + usable_runtime_by_id[result.id].append(result.runtime) - :return: The runtime in nanoseconds. - """ - return sum( - [ - min(usable_runtime_data) - for invocation_id, usable_runtime_data in self.usable_runtime_data_by_test_case().items() - ] - ) + return usable_runtime_by_id - def usable_replay_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: - """Collect runtime data for replay tests that passed and have runtime information. - :return: A dictionary mapping invocation IDs to lists of runtime values. - """ - usable_runtimes = [ - (result.id, result.runtime) - for result in self.test_results - if result.did_pass and result.runtime and result.test_type == TestType.REPLAY_TEST - ] - - return { - usable_id: [runtime[1] for runtime in usable_runtimes if runtime[0] == usable_id] - for usable_id in {runtime[0] for runtime in usable_runtimes} - } + def total_passed_runtime(self) -> int: + """Calculate the sum of runtimes of all test cases that passed. - def total_replay_test_runtime(self) -> int: - """Calculate the sum of runtimes of replay test cases that passed, where a testcase runtime - is the minimum value of all looped execution runtimes. + A testcase runtime is the minimum value of all looped execution runtimes. :return: The runtime in nanoseconds. """ - replay_runtime_data = self.usable_replay_runtime_data_by_test_case() - - return sum([ - min(runtimes) - for invocation_id, runtimes in replay_runtime_data.items() - ]) if replay_runtime_data else 0 + return sum( + [min(usable_runtime_data) for _, usable_runtime_data in self.usable_runtime_data_by_test_case().items()] + ) def __iter__(self) -> Iterator[FunctionTestInvocation]: return iter(self.test_results) diff --git a/tests/test_instrument_codeflash_trace.py b/tests/test_instrument_codeflash_trace.py index 967d5d6f0..6b884c631 100644 --- a/tests/test_instrument_codeflash_trace.py +++ b/tests/test_instrument_codeflash_trace.py @@ -1,9 +1,10 @@ from __future__ import annotations +import tempfile from pathlib import Path -from codeflash.benchmarking.instrument_codeflash_trace import add_codeflash_decorator_to_code - +from codeflash.benchmarking.instrument_codeflash_trace import add_codeflash_decorator_to_code, \ + instrument_codeflash_trace_decorator from codeflash.discovery.functions_to_optimize import FunctionParent, FunctionToOptimize @@ -22,7 +23,7 @@ def normal_function(): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -34,6 +35,7 @@ def normal_function(): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_normal_method() -> None: """Test adding decorator to a normal method.""" code = """ @@ -50,7 +52,7 @@ def normal_method(self): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -63,6 +65,7 @@ def normal_method(self): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_classmethod() -> None: """Test adding decorator to a classmethod.""" code = """ @@ -80,7 +83,7 @@ def class_method(cls): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -94,6 +97,7 @@ def class_method(cls): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_staticmethod() -> None: """Test adding decorator to a staticmethod.""" code = """ @@ -111,7 +115,7 @@ def static_method(): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -125,6 +129,7 @@ def static_method(): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_init_function() -> None: """Test adding decorator to an __init__ function.""" code = """ @@ -141,7 +146,7 @@ def __init__(self, value): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -154,6 +159,7 @@ def __init__(self, value): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_with_multiple_decorators() -> None: """Test adding decorator to a function with multiple existing decorators.""" code = """ @@ -172,7 +178,7 @@ def property_method(self): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -187,6 +193,7 @@ def property_method(self): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_function_in_multiple_classes() -> None: """Test that only the right class's method gets the decorator.""" code = """ @@ -207,7 +214,7 @@ def test_method(self): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) expected_code = """ @@ -224,6 +231,7 @@ def test_method(self): assert modified_code.strip() == expected_code.strip() + def test_add_decorator_to_nonexistent_function() -> None: """Test that code remains unchanged when function doesn't exist.""" code = """ @@ -239,8 +247,223 @@ def existing_function(): modified_code = add_codeflash_decorator_to_code( code=code, - function_to_optimize=fto + functions_to_optimize=[fto] ) # Code should remain unchanged assert modified_code.strip() == code.strip() + + +def test_add_decorator_to_multiple_functions() -> None: + """Test adding decorator to multiple functions.""" + code = """ +def function_one(): + return "First function" + +class TestClass: + def method_one(self): + return "First method" + + def method_two(self): + return "Second method" + +def function_two(): + return "Second function" +""" + + functions_to_optimize = [ + FunctionToOptimize( + function_name="function_one", + file_path=Path("dummy_path.py"), + parents=[] + ), + FunctionToOptimize( + function_name="method_two", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ), + FunctionToOptimize( + function_name="function_two", + file_path=Path("dummy_path.py"), + parents=[] + ) + ] + + modified_code = add_codeflash_decorator_to_code( + code=code, + functions_to_optimize=functions_to_optimize + ) + + expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +@codeflash_trace +def function_one(): + return "First function" + +class TestClass: + def method_one(self): + return "First method" + + @codeflash_trace + def method_two(self): + return "Second method" + +@codeflash_trace +def function_two(): + return "Second function" +""" + + assert modified_code.strip() == expected_code.strip() + + +def test_instrument_codeflash_trace_decorator_single_file() -> None: + """Test instrumenting codeflash trace decorator on a single file.""" + with tempfile.TemporaryDirectory() as temp_dir: + # Create a test Python file + test_file_path = Path(temp_dir) / "test_module.py" + test_file_content = """ +def function_one(): + return "First function" + +class TestClass: + def method_one(self): + return "First method" + + def method_two(self): + return "Second method" + +def function_two(): + return "Second function" +""" + test_file_path.write_text(test_file_content, encoding="utf-8") + + # Define functions to optimize + functions_to_optimize = [ + FunctionToOptimize( + function_name="function_one", + file_path=test_file_path, + parents=[] + ), + FunctionToOptimize( + function_name="method_two", + file_path=test_file_path, + parents=[FunctionParent(name="TestClass", type="ClassDef")] + ) + ] + + # Execute the function being tested + instrument_codeflash_trace_decorator({test_file_path: functions_to_optimize}) + + # Read the modified file + modified_content = test_file_path.read_text(encoding="utf-8") + + # Define expected content (with isort applied) + expected_content = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace + + +@codeflash_trace +def function_one(): + return "First function" + +class TestClass: + def method_one(self): + return "First method" + + @codeflash_trace + def method_two(self): + return "Second method" + +def function_two(): + return "Second function" +""" + + # Compare the modified content with expected content + assert modified_content.strip() == expected_content.strip() + + +def test_instrument_codeflash_trace_decorator_multiple_files() -> None: + """Test instrumenting codeflash trace decorator on multiple files.""" + with tempfile.TemporaryDirectory() as temp_dir: + # Create first test Python file + test_file_1_path = Path(temp_dir) / "module_a.py" + test_file_1_content = """ +def function_a(): + return "Function in module A" + +class ClassA: + def method_a(self): + return "Method in ClassA" +""" + test_file_1_path.write_text(test_file_1_content, encoding="utf-8") + + # Create second test Python file + test_file_2_path = Path(temp_dir) / "module_b.py" + test_file_2_content =""" +def function_b(): + return "Function in module B" + +class ClassB: + @staticmethod + def static_method_b(): + return "Static method in ClassB" +""" + test_file_2_path.write_text(test_file_2_content, encoding="utf-8") + + # Define functions to optimize + file_to_funcs_to_optimize = { + test_file_1_path: [ + FunctionToOptimize( + function_name="function_a", + file_path=test_file_1_path, + parents=[] + ) + ], + test_file_2_path: [ + FunctionToOptimize( + function_name="static_method_b", + file_path=test_file_2_path, + parents=[FunctionParent(name="ClassB", type="ClassDef")] + ) + ] + } + + # Execute the function being tested + instrument_codeflash_trace_decorator(file_to_funcs_to_optimize) + + # Read the modified files + modified_content_1 = test_file_1_path.read_text(encoding="utf-8") + modified_content_2 = test_file_2_path.read_text(encoding="utf-8") + + # Define expected content for first file (with isort applied) + expected_content_1 = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace + + +@codeflash_trace +def function_a(): + return "Function in module A" + +class ClassA: + def method_a(self): + return "Method in ClassA" +""" + + # Define expected content for second file (with isort applied) + expected_content_2 = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace + + +def function_b(): + return "Function in module B" + +class ClassB: + @staticmethod + @codeflash_trace + def static_method_b(): + return "Static method in ClassB" +""" + + # Compare the modified content with expected content + assert modified_content_1.strip() == expected_content_1.strip() + assert modified_content_2.strip() == expected_content_2.strip() \ No newline at end of file diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index fcc5b0f67..f08b2485a 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,6 +1,6 @@ import sqlite3 -from codeflash.benchmarking.benchmark_database_utils import BenchmarkDatabaseUtils +from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.replay_test import generate_replay_test from pathlib import Path @@ -27,7 +27,7 @@ def test_trace_benchmarks(): # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM function_calls ORDER BY benchmark_file_name, benchmark_function_name, function_name") + "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_name, benchmark_function_name, function_name") function_calls = cursor.fetchall() # Assert the length of function calls @@ -154,7 +154,6 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): finally: # cleanup shutil.rmtree(tests_root) - pass def test_trace_multithreaded_benchmark() -> None: project_root = Path(__file__).parent.parent / "code_to_optimize" @@ -173,13 +172,13 @@ def test_trace_multithreaded_benchmark() -> None: # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM function_calls ORDER BY benchmark_file_name, benchmark_function_name, function_name") + "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_name, benchmark_function_name, function_name") function_calls = cursor.fetchall() # Assert the length of function calls assert len(function_calls) == 10, f"Expected 10 function calls, but got {len(function_calls)}" - function_benchmark_timings = BenchmarkDatabaseUtils.get_function_benchmark_timings(output_file) - total_benchmark_timings = BenchmarkDatabaseUtils.get_benchmark_timings(output_file) + 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 @@ -208,5 +207,4 @@ def test_trace_multithreaded_benchmark() -> None: finally: # cleanup - shutil.rmtree(tests_root) - pass \ No newline at end of file + shutil.rmtree(tests_root) \ No newline at end of file From 582bea066c502d0df7f4b09142a7b59b4d194805 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 28 Mar 2025 14:58:49 -0700 Subject: [PATCH 084/122] started implementing group by benchmark --- codeflash/benchmarking/plugin/plugin.py | 11 ++---- codeflash/benchmarking/replay_test.py | 5 +-- codeflash/benchmarking/utils.py | 2 +- codeflash/models/models.py | 5 +-- codeflash/optimization/function_optimizer.py | 40 +++++++++++--------- codeflash/optimization/optimizer.py | 9 +++-- codeflash/verification/test_results.py | 20 +++++++++- 7 files changed, 56 insertions(+), 36 deletions(-) diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index 9d7da6ef2..a4805cca3 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -101,8 +101,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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func) # Initialize the inner dictionary if needed if qualified_name not in result: result[qualified_name] = {} @@ -152,8 +151,7 @@ def get_benchmark_timings(trace_path: Path) -> dict[BenchmarkKey, int]: # Process overhead information for row in cursor.fetchall(): benchmark_file, benchmark_func, benchmark_line, total_overhead_ns = row - benchmark_key = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + benchmark_key = BenchmarkKey(file_name=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 @@ -167,8 +165,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 = f"{benchmark_file}::{benchmark_func}::{benchmark_line}" - benchmark_key = BenchmarkKey(file_name=benchmark_file, function_name=benchmark_func, line_number=benchmark_line) + benchmark_key = BenchmarkKey(file_name=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 @@ -239,7 +236,7 @@ def test_something(benchmark): The return value of the function """ - benchmark_file_name = self.request.node.fspath.basename + benchmark_file_name = self.request.node.fspath benchmark_function_name = self.request.node.name line_number = int(str(sys._getframe(1).f_lineno)) # 1 frame up in the call stack diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 670d6e4bd..58bae35f1 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -196,12 +196,11 @@ def create_trace_replay_test_code( return imports + "\n" + metadata + "\n" + test_template def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework: str = "pytest", max_run_count: int = 100) -> int: - """Generate multiple replay tests from the traced function calls, grouping by benchmark name. + """Generate multiple replay tests from the traced function calls, grouped by benchmark. Args: trace_file_path: Path to the SQLite database file output_dir: Directory to write the generated tests (if None, only returns the code) - project_root: Root directory of the project for module imports test_framework: 'pytest' or 'unittest' max_run_count: Maximum number of runs to include per function @@ -267,7 +266,7 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework # Write to file if requested if output_dir: output_file = get_test_file_path( - test_dir=Path(output_dir), function_name=f"{benchmark_file_name[5:]}_{benchmark_function_name}", test_type="replay" + test_dir=Path(output_dir), function_name=f"{benchmark_file_name}_{benchmark_function_name}", test_type="replay" ) # Write test code to file, parents = true output_dir.mkdir(parents=True, exist_ok=True) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 5f14f141f..feb9ed0fc 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -93,7 +93,7 @@ def process_benchmark_data( for benchmark_key, og_benchmark_timing in fto_benchmark_timings.items(): try: - benchmark_file_name, benchmark_test_function, line_number = benchmark_key.split("::") + benchmark_file_name, benchmark_test_function = benchmark_key.split("::") except ValueError: continue # Skip malformed benchmark keys diff --git a/codeflash/models/models.py b/codeflash/models/models.py index e046cf910..f62131e2a 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -82,10 +82,9 @@ class BestOptimization(BaseModel): class BenchmarkKey: file_name: str function_name: str - line_number: int def __str__(self) -> str: - return f"{self.file_name}::{self.function_name}::{self.line_number}" + return f"{self.file_name}::{self.function_name}" @dataclass class BenchmarkDetail: @@ -270,7 +269,7 @@ class FunctionParent: class OriginalCodeBaseline(BaseModel): behavioral_test_results: TestResults benchmarking_test_results: TestResults - replay_benchmarking_test_results: Optional[TestResults] = None + replay_benchmarking_test_results: Optional[dict[BenchmarkKey, TestResults]] = None runtime: int coverage_results: Optional[CoverageData] diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 464027778..c0143eb10 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -93,8 +93,8 @@ def __init__( function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_ast: ast.FunctionDef | None = None, aiservice_client: AiServiceClient | None = None, - function_benchmark_timings: dict[str, int] | None = None, - total_benchmark_timings: dict[str, int] | None = None, + function_benchmark_timings: dict[BenchmarkKey, int] | None = None, + total_benchmark_timings: dict[BenchmarkKey, int] | None = None, args: Namespace | None = None, ) -> None: self.project_root = test_cfg.project_root_path @@ -433,20 +433,24 @@ def determine_best_candidate( tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") if self.args.benchmark: - original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results.total_passed_runtime() - candidate_replay_runtime = candidate_result.replay_benchmarking_test_results.total_passed_runtime() - replay_perf_gain = performance_gain( - original_runtime_ns=original_code_replay_runtime, - optimized_runtime_ns=candidate_replay_runtime, - ) - tree.add(f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") - tree.add( - f"Best benchmark replay runtime: {humanize_runtime(candidate_replay_runtime)} " - f"(measured over {candidate_result.max_loop_count} " - f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" - ) - tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain + 1:.1f}X") + + benchmark_keys = {(benchmark.file_name, benchmark.function_name) for benchmark in self.total_benchmark_timings} + test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmark(benchmark_keys) + for benchmark_key, test_results in test_results_by_benchmark.items(): + original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results[benchmark_key].total_passed_runtime() + candidate_replay_runtime = candidate_result.replay_benchmarking_test_results.total_passed_runtime() + replay_perf_gain = performance_gain( + original_runtime_ns=original_code_replay_runtime, + optimized_runtime_ns=candidate_replay_runtime, + ) + tree.add(f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") + tree.add( + f"Best benchmark replay runtime: {humanize_runtime(candidate_replay_runtime)} " + f"(measured over {candidate_result.max_loop_count} " + f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + ) + tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain + 1:.1f}X") best_optimization = BestOptimization( candidate=candidate, helper_functions=code_context.helper_functions, @@ -949,7 +953,7 @@ def establish_original_code_baseline( logger.debug(f"Total original code runtime (ns): {total_timing}") if self.args.benchmark: - replay_benchmarking_test_results = benchmarking_results.filter(TestType.REPLAY_TEST) + replay_benchmarking_test_results = benchmarking_results.filter_by_test_type(TestType.REPLAY_TEST) logger.info(f"Total replay test runtime: {humanize_runtime(replay_benchmarking_test_results.total_passed_runtime())}") return Success( ( @@ -1071,7 +1075,7 @@ def run_optimized_candidate( logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") if self.args.benchmark: - candidate_replay_benchmarking_results = candidate_benchmarking_results.filter(TestType.REPLAY_TEST) + candidate_replay_benchmarking_results = candidate_benchmarking_results.filter_by_test_type(TestType.REPLAY_TEST) logger.debug( f"Total optimized code {optimization_candidate_index} replay benchmark runtime (ns): {candidate_replay_benchmarking_results.total_passed_runtime()}" ) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 4d17a5255..35d91a274 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -60,8 +60,8 @@ def create_function_optimizer( function_to_optimize_ast: ast.FunctionDef | None = None, function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, function_to_optimize_source_code: str | None = "", - function_benchmark_timings: dict[str, dict[str, float]] | None = None, - total_benchmark_timings: dict[str, float] | None = None, + function_benchmark_timings: dict[str, dict[BenchmarkKey, float]] | None = None, + total_benchmark_timings: dict[BenchmarkKey, float] | None = None, ) -> FunctionOptimizer: return FunctionOptimizer( function_to_optimize=function_to_optimize, @@ -111,7 +111,10 @@ def run(self) -> None: try: instrument_codeflash_trace_decorator(file_to_funcs_to_optimize) trace_file = Path(self.args.benchmarks_root) / "benchmarks.trace" - replay_tests_dir = Path(self.args.tests_root) / "codeflash_replay_tests" + if trace_file.exists(): + trace_file.unlink() + + replay_tests_dir = Path(self.args.tests_root) trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Run all tests that use pytest-benchmark replay_count = generate_replay_test(trace_file, replay_tests_dir) if replay_count == 0: diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py index 6f4f397ab..25f258e26 100644 --- a/codeflash/verification/test_results.py +++ b/codeflash/verification/test_results.py @@ -125,7 +125,7 @@ def merge(self, other: TestResults) -> None: raise ValueError(msg) self.test_result_idx[k] = v + original_len - def filter(self, test_type: TestType) -> TestResults: + def filter_by_test_type(self, test_type: TestType) -> TestResults: filtered_test_results = [] filtered_test_results_idx = {} for test_result in self.test_results: @@ -134,6 +134,24 @@ def filter(self, test_type: TestType) -> TestResults: filtered_test_results.append(test_result) return TestResults(test_results=filtered_test_results, test_result_idx=filtered_test_results_idx) + def group_by_benchmark(self, benchmark_key_set:set[tuple[str,str]]) -> dict[tuple[str,str],TestResults]: + """Group TestResults by benchmark key. + + For now, use a tuple of (file_path, function_name) as the benchmark key. Can't import BenchmarkKey because of circular import. + + Args: + benchmark_key_set (set[tuple[str,str]]): A set of tuples of (file_path, function_name) + + Returns: + TestResults: A new TestResults object with the test results grouped by benchmark key. + + """ + test_result_by_benchmark = defaultdict(TestResults) + for test_result in self.test_results: + if test_result.test_type == TestType.REPLAY_TEST and (test_result.id.test_module_path,test_result.id.test_function_name) in benchmark_key_set: + test_result_by_benchmark[(test_result.id.test_module_path,test_result.id.test_function_name)].add(test_result) + return test_result_by_benchmark + def get_by_unique_invocation_loop_id(self, unique_invocation_loop_id: str) -> FunctionTestInvocation | None: try: return self.test_results[self.test_result_idx[unique_invocation_loop_id]] From e5a8260ac982b6ae41d22f37c61ae21f8a646b99 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Mon, 31 Mar 2025 16:47:58 -0700 Subject: [PATCH 085/122] reworked matching benchmark key to test results. --- codeflash/benchmarking/codeflash_trace.py | 12 +- codeflash/benchmarking/plugin/plugin.py | 24 +- codeflash/benchmarking/replay_test.py | 64 +-- codeflash/benchmarking/utils.py | 32 +- codeflash/github/PrComment.py | 2 +- codeflash/models/models.py | 486 ++++++++++--------- codeflash/optimization/function_optimizer.py | 118 ++--- codeflash/optimization/optimizer.py | 7 +- codeflash/result/explanation.py | 3 +- tests/test_trace_benchmarks.py | 42 +- tests/test_unit_test_discovery.py | 1 - 11 files changed, 391 insertions(+), 400 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 2ae57307b..8c307f8a0 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -33,8 +33,8 @@ def setup(self, trace_path: str) -> None: cur.execute("PRAGMA synchronous = OFF") cur.execute( "CREATE TABLE IF NOT EXISTS benchmark_function_timings(" - "function_name TEXT, class_name TEXT, module_name TEXT, file_name TEXT," - "benchmark_function_name TEXT, benchmark_file_name TEXT, benchmark_line_number INTEGER," + "function_name TEXT, class_name TEXT, module_name TEXT, file_path TEXT," + "benchmark_function_name TEXT, benchmark_file_path TEXT, benchmark_line_number INTEGER," "function_time_ns INTEGER, overhead_time_ns INTEGER, args BLOB, kwargs BLOB)" ) self._connection.commit() @@ -62,8 +62,8 @@ def write_function_timings(self) -> None: # Insert data into the benchmark_function_timings table cur.executemany( "INSERT INTO benchmark_function_timings" - "(function_name, class_name, module_name, file_name, benchmark_function_name, " - "benchmark_file_name, benchmark_line_number, function_time_ns, overhead_time_ns, args, kwargs) " + "(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) " "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", self.function_calls_data ) @@ -115,7 +115,7 @@ def wrapper(*args, **kwargs): # 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", "") + benchmark_file_path = os.environ.get("CODEFLASH_BENCHMARK_FILE_PATH", "") benchmark_line_number = os.environ.get("CODEFLASH_BENCHMARK_LINE_NUMBER", "") # Get class name class_name = "" @@ -151,7 +151,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_name, benchmark_line_number, execution_time, + benchmark_function_name, benchmark_file_path, benchmark_line_number, execution_time, overhead_time, pickled_args, pickled_kwargs) ) return result diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index a4805cca3..09858601c 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -24,7 +24,7 @@ def setup(self, trace_path:str) -> None: cur.execute("PRAGMA synchronous = OFF") cur.execute( "CREATE TABLE IF NOT EXISTS benchmark_timings(" - "benchmark_file_name TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER," + "benchmark_file_path TEXT, benchmark_function_name TEXT, benchmark_line_number INTEGER," "benchmark_time_ns INTEGER)" ) self._connection.commit() @@ -47,7 +47,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_name, benchmark_function_name, benchmark_line_number, benchmark_time_ns) VALUES (?, ?, ?, ?)", + "INSERT INTO benchmark_timings (benchmark_file_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns) VALUES (?, ?, ?, ?)", self.benchmark_timings ) self._connection.commit() @@ -86,7 +86,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_name, benchmark_function_name, benchmark_line_number, function_time_ns " + "benchmark_file_path, benchmark_function_name, benchmark_line_number, function_time_ns " "FROM benchmark_function_timings" ) @@ -101,7 +101,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_name=benchmark_file, function_name=benchmark_func) + benchmark_key = BenchmarkKey(file_path=benchmark_file, function_name=benchmark_func) # Initialize the inner dictionary if needed if qualified_name not in result: result[qualified_name] = {} @@ -143,20 +143,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_name, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " + "SELECT benchmark_file_path, benchmark_function_name, benchmark_line_number, SUM(overhead_time_ns) " "FROM benchmark_function_timings " - "GROUP BY benchmark_file_name, benchmark_function_name, benchmark_line_number" + "GROUP BY benchmark_file_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_name=benchmark_file, function_name=benchmark_func) + benchmark_key = BenchmarkKey(file_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_name, benchmark_function_name, benchmark_line_number, benchmark_time_ns " + "SELECT benchmark_file_path, benchmark_function_name, benchmark_line_number, benchmark_time_ns " "FROM benchmark_timings" ) @@ -165,7 +165,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_name=benchmark_file, function_name=benchmark_func) + benchmark_key = BenchmarkKey(file_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 @@ -236,13 +236,13 @@ def test_something(benchmark): The return value of the function """ - benchmark_file_name = self.request.node.fspath + benchmark_file_path = str(self.request.node.fspath) 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_NAME"] = benchmark_file_name + os.environ["CODEFLASH_BENCHMARK_FILE_PATH"] = benchmark_file_path os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(line_number) os.environ["CODEFLASH_BENCHMARKING"] = "True" @@ -260,7 +260,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_name, benchmark_function_name, line_number, end - start)) + (benchmark_file_path, benchmark_function_name, line_number, end - start)) return result diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 58bae35f1..9ecac2ec4 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -2,18 +2,21 @@ import sqlite3 import textwrap -from collections.abc import Generator -from typing import Any, Dict +from pathlib import Path +from typing import TYPE_CHECKING, Any import isort from codeflash.cli_cmds.console import logger from codeflash.discovery.functions_to_optimize import inspect_top_level_functions_or_methods from codeflash.verification.verification_utils import get_test_file_path -from pathlib import Path + +if TYPE_CHECKING: + from collections.abc import Generator + def get_next_arg_and_return( - trace_file: str, function_name: str, file_name: str, class_name: str | None = None, num_to_get: int = 256 + trace_file: str, function_name: str, file_path: str, class_name: str | None = None, num_to_get: int = 256 ) -> Generator[Any]: db = sqlite3.connect(trace_file) cur = db.cursor() @@ -21,13 +24,13 @@ def get_next_arg_and_return( if class_name is not None: cursor = cur.execute( - "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_name = ? AND class_name = ? LIMIT ?", - (function_name, file_name, class_name, limit), + "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_path = ? AND class_name = ? LIMIT ?", + (function_name, file_path, class_name, limit), ) else: cursor = cur.execute( - "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_name = ? AND class_name = '' LIMIT ?", - (function_name, file_name, limit), + "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_path = ? AND class_name = '' LIMIT ?", + (function_name, file_path, limit), ) while (val := cursor.fetchone()) is not None: @@ -88,7 +91,7 @@ def create_trace_replay_test_code( # Templates for different types of tests test_function_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = {function_name}(*args, **kwargs) @@ -97,7 +100,7 @@ def create_trace_replay_test_code( test_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} function_name = "{orig_function_name}" @@ -112,7 +115,7 @@ def create_trace_replay_test_code( test_class_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} if not args: @@ -122,7 +125,7 @@ def create_trace_replay_test_code( ) test_static_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_name=r"{file_name}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} ret = {class_name_alias}{method_name}(*args, **kwargs) @@ -140,13 +143,13 @@ def create_trace_replay_test_code( module_name = func.get("module_name") function_name = func.get("function_name") class_name = func.get("class_name") - file_name = func.get("file_name") + file_path = func.get("file_path") function_properties = func.get("function_properties") if not class_name: alias = get_function_alias(module_name, function_name) test_body = test_function_body.format( function_name=alias, - file_name=file_name, + file_path=file_path, orig_function_name=function_name, max_run_count=max_run_count, ) @@ -160,7 +163,7 @@ def create_trace_replay_test_code( if function_properties.is_classmethod: test_body = test_class_method_body.format( orig_function_name=function_name, - file_name=file_name, + file_path=file_path, class_name_alias=class_name_alias, class_name=class_name, method_name=method_name, @@ -170,7 +173,7 @@ def create_trace_replay_test_code( elif function_properties.is_staticmethod: test_body = test_static_method_body.format( orig_function_name=function_name, - file_name=file_name, + file_path=file_path, class_name_alias=class_name_alias, class_name=class_name, method_name=method_name, @@ -180,7 +183,7 @@ def create_trace_replay_test_code( else: test_body = test_method_body.format( orig_function_name=function_name, - file_name=file_name, + file_path=file_path, class_name_alias=class_name_alias, class_name=class_name, method_name=method_name, @@ -216,42 +219,41 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework # Get distinct benchmark names cursor.execute( - "SELECT DISTINCT benchmark_function_name, benchmark_file_name FROM benchmark_function_timings" + "SELECT DISTINCT benchmark_function_name, benchmark_file_path FROM benchmark_function_timings" ) benchmarks = cursor.fetchall() # Generate a test for each benchmark for benchmark in benchmarks: - benchmark_function_name, benchmark_file_name = benchmark + benchmark_function_name, benchmark_file_path = benchmark # Get functions associated with this benchmark cursor.execute( - "SELECT DISTINCT function_name, class_name, module_name, file_name, benchmark_line_number FROM benchmark_function_timings " - "WHERE benchmark_function_name = ? AND benchmark_file_name = ?", - (benchmark_function_name, benchmark_file_name) + "SELECT DISTINCT function_name, class_name, module_name, file_path, benchmark_line_number FROM benchmark_function_timings " + "WHERE benchmark_function_name = ? AND benchmark_file_path = ?", + (benchmark_function_name, benchmark_file_path) ) functions_data = [] for func_row in cursor.fetchall(): - function_name, class_name, module_name, file_name, benchmark_line_number = func_row - + function_name, class_name, module_name, file_path, benchmark_line_number = func_row # Add this function to our list functions_data.append({ "function_name": function_name, "class_name": class_name, - "file_name": file_name, + "file_path": file_path, "module_name": module_name, "benchmark_function_name": benchmark_function_name, - "benchmark_file_name": benchmark_file_name, + "benchmark_file_path": benchmark_file_path, "benchmark_line_number": benchmark_line_number, "function_properties": inspect_top_level_functions_or_methods( - file_name=file_name, + file_name=Path(file_path), function_or_method_name=function_name, class_name=class_name, ) }) if not functions_data: - logger.info(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_name}") + logger.info(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_path}") continue # Generate the test code for this benchmark @@ -265,17 +267,19 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework # Write to file if requested if output_dir: + name = Path(benchmark_file_path).name.split(".")[0][5:] # remove "test_" from the name since we add it in later output_file = get_test_file_path( - test_dir=Path(output_dir), function_name=f"{benchmark_file_name}_{benchmark_function_name}", test_type="replay" + test_dir=Path(output_dir), function_name=f"{name}_{benchmark_function_name}", test_type="replay" ) # Write test code to file, parents = true output_dir.mkdir(parents=True, exist_ok=True) output_file.write_text(test_code, "utf-8") count += 1 - logger.info(f"Replay test for benchmark `{benchmark_function_name}` in {benchmark_file_name} written to {output_file}") + logger.info(f"Replay test for benchmark `{benchmark_function_name}` in {name} written to {output_file}") conn.close() except Exception as e: logger.info(f"Error generating replay tests: {e}") + return count \ No newline at end of file diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index feb9ed0fc..1d8b22f50 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -11,7 +11,7 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[BenchmarkKey, int]], - total_benchmark_timings: dict[BenchmarkKey, int]) -> dict[str, list[tuple[str, float, float, float]]]: + total_benchmark_timings: dict[BenchmarkKey, int]) -> dict[str, list[tuple[BenchmarkKey, float, float, float]]]: function_to_result = {} # Process each function's benchmark data for func_path, test_times in function_benchmark_timings.items(): @@ -23,18 +23,18 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, di logger.debug(f"Skipping test {benchmark_key} due to func_time {func_time} > total_time {total_time}") # If the function time is greater than total time, likely to have multithreading / multiprocessing issues. # Do not try to project the optimization impact for this function. - sorted_tests.append((str(benchmark_key), 0.0, 0.0, 0.0)) + sorted_tests.append((benchmark_key, 0.0, 0.0, 0.0)) if total_time > 0: percentage = (func_time / total_time) * 100 # Convert nanoseconds to milliseconds func_time_ms = func_time / 1_000_000 total_time_ms = total_time / 1_000_000 - sorted_tests.append((str(benchmark_key), total_time_ms, func_time_ms, percentage)) + sorted_tests.append((benchmark_key, total_time_ms, func_time_ms, percentage)) sorted_tests.sort(key=lambda x: x[3], reverse=True) function_to_result[func_path] = sorted_tests return function_to_result -def print_benchmark_table(function_to_results: dict[str, list[tuple[str, float, float, float]]]) -> None: +def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey, float, float, float]]]) -> None: console = Console() for func_path, sorted_tests in function_to_results.items(): function_name = func_path.split(":")[-1] @@ -48,19 +48,17 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[str, float, table.add_column("Function Time (ms)", justify="right", style="yellow") table.add_column("Percentage (%)", justify="right", style="red") - for test_name, total_time, func_time, percentage in sorted_tests: - benchmark_file, benchmark_func, benchmark_line = test_name.split("::") - benchmark_name = f"{benchmark_file}::{benchmark_func}" + for benchmark_key, total_time, func_time, percentage in sorted_tests: if total_time == 0.0: table.add_row( - benchmark_name, + f"{benchmark_key.file_path}::{benchmark_key.function_name}", "N/A", "N/A", "N/A" ) else: table.add_row( - benchmark_name, + f"{benchmark_key.file_path}::{benchmark_key.function_name}", f"{total_time:.3f}", f"{func_time:.3f}", f"{percentage:.2f}" @@ -71,9 +69,9 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[str, float, def process_benchmark_data( - replay_performance_gain: float, - fto_benchmark_timings: dict[str, int], - total_benchmark_timings: dict[str, int] + replay_performance_gain: dict[BenchmarkKey, float], + fto_benchmark_timings: dict[BenchmarkKey, int], + total_benchmark_timings: dict[BenchmarkKey, int] ) -> Optional[ProcessedBenchmarkInfo]: """Process benchmark data and generate detailed benchmark information. @@ -92,10 +90,6 @@ def process_benchmark_data( benchmark_details = [] for benchmark_key, og_benchmark_timing in fto_benchmark_timings.items(): - try: - benchmark_file_name, benchmark_test_function = benchmark_key.split("::") - except ValueError: - continue # Skip malformed benchmark keys total_benchmark_timing = total_benchmark_timings.get(benchmark_key, 0) @@ -104,7 +98,7 @@ def process_benchmark_data( # Calculate expected new benchmark timing expected_new_benchmark_timing = total_benchmark_timing - og_benchmark_timing + ( - 1 / (replay_performance_gain + 1) + 1 / (replay_performance_gain[benchmark_key] + 1) ) * og_benchmark_timing # Calculate speedup @@ -112,8 +106,8 @@ def process_benchmark_data( benchmark_details.append( BenchmarkDetail( - benchmark_name=benchmark_file_name, - test_function=benchmark_test_function, + benchmark_name=benchmark_key.file_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)), speedup_percent=benchmark_speedup_percent diff --git a/codeflash/github/PrComment.py b/codeflash/github/PrComment.py index 5b891b8a5..1e66c5608 100644 --- a/codeflash/github/PrComment.py +++ b/codeflash/github/PrComment.py @@ -6,7 +6,7 @@ from codeflash.code_utils.time_utils import humanize_runtime from codeflash.models.models import BenchmarkDetail -from codeflash.verification.test_results import TestResults +from codeflash.models.models import TestResults @dataclass(frozen=True, config={"arbitrary_types_allowed": True}) diff --git a/codeflash/models/models.py b/codeflash/models/models.py index f62131e2a..ed0360eef 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -1,30 +1,31 @@ from __future__ import annotations +from collections import defaultdict +from typing import TYPE_CHECKING + +from rich.tree import Tree + +from codeflash.cli_cmds.console import DEBUG_MODE + +if TYPE_CHECKING: + from collections.abc import Iterator import enum -import json import re -from collections.abc import Collection, Iterator +import sys +from collections.abc import Collection from enum import Enum, IntEnum from pathlib import Path from re import Pattern -from typing import Annotated, Any, Optional, Union +from typing import Annotated, Any, Optional, Union, cast -import sentry_sdk -from coverage.exceptions import NoDataError from jedi.api.classes import Name from pydantic import AfterValidator, BaseModel, ConfigDict, Field from pydantic.dataclasses import dataclass from codeflash.cli_cmds.console import console, logger -from codeflash.code_utils.code_utils import validate_python_code -from codeflash.code_utils.coverage_utils import ( - build_fully_qualified_name, - extract_dependent_function, - generate_candidates, -) +from codeflash.code_utils.code_utils import module_name_from_file_path, validate_python_code from codeflash.code_utils.env_utils import is_end_to_end -from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.verification.test_results import TestResults, TestType +from codeflash.verification.comparator import comparator # If the method spam is in the class Ham, which is at the top level of the module eggs in the package foo, the fully # qualified name of the method is foo.eggs.Ham.spam, its qualified name is Ham.spam, and its name is spam. The full name @@ -72,19 +73,18 @@ class BestOptimization(BaseModel): candidate: OptimizedCandidate helper_functions: list[FunctionSource] runtime: int - replay_runtime: Optional[int] = None - replay_performance_gain: Optional[float] = None + replay_performance_gain: Optional[dict[BenchmarkKey,float]] = None winning_behavioral_test_results: TestResults winning_benchmarking_test_results: TestResults winning_replay_benchmarking_test_results : Optional[TestResults] = None @dataclass(frozen=True) class BenchmarkKey: - file_name: str + file_path: str function_name: str def __str__(self) -> str: - return f"{self.file_name}::{self.function_name}" + return f"{self.file_path}::{self.function_name}" @dataclass class BenchmarkDetail: @@ -164,7 +164,7 @@ class OptimizedCandidateResult(BaseModel): best_test_runtime: int behavior_test_results: TestResults benchmarking_test_results: TestResults - replay_benchmarking_test_results: Optional[TestResults] = None + replay_benchmarking_test_results: Optional[dict[BenchmarkKey, TestResults]] = None optimization_candidate_index: int total_candidate_timing: int @@ -295,209 +295,6 @@ class CoverageData: blank_re: Pattern[str] = re.compile(r"\s*(#|$)") else_re: Pattern[str] = re.compile(r"\s*else\s*:\s*(#|$)") - @staticmethod - def load_from_sqlite_database( - database_path: Path, config_path: Path, function_name: str, code_context: CodeOptimizationContext, source_code_path: Path - ) -> CoverageData: - """Load coverage data from an SQLite database, mimicking the behavior of load_from_coverage_file.""" - from coverage import Coverage - from coverage.jsonreport import JsonReporter - - cov = Coverage(data_file=database_path,config_file=config_path, data_suffix=True, auto_data=True, branch=True) - - if not database_path.stat().st_size or not database_path.exists(): - logger.debug(f"Coverage database {database_path} is empty or does not exist") - sentry_sdk.capture_message(f"Coverage database {database_path} is empty or does not exist") - return CoverageData.create_empty(source_code_path, function_name, code_context) - cov.load() - - reporter = JsonReporter(cov) - temp_json_file = database_path.with_suffix(".report.json") - with temp_json_file.open("w") as f: - try: - reporter.report(morfs=[source_code_path.as_posix()], outfile=f) - except NoDataError: - sentry_sdk.capture_message(f"No coverage data found for {function_name} in {source_code_path}") - return CoverageData.create_empty(source_code_path, function_name, code_context) - with temp_json_file.open() as f: - original_coverage_data = json.load(f) - - coverage_data, status = CoverageData._parse_coverage_file(temp_json_file, source_code_path) - - main_func_coverage, dependent_func_coverage = CoverageData._fetch_function_coverages( - function_name, code_context, coverage_data, original_cov_data=original_coverage_data - ) - - total_executed_lines, total_unexecuted_lines = CoverageData._aggregate_coverage( - main_func_coverage, dependent_func_coverage - ) - - total_lines = total_executed_lines | total_unexecuted_lines - coverage = len(total_executed_lines) / len(total_lines) * 100 if total_lines else 0.0 - # coverage = (lines covered of the original function + its 1 level deep helpers) / (lines spanned by original function + its 1 level deep helpers), if no helpers then just the original function coverage - - functions_being_tested = [main_func_coverage.name] - if dependent_func_coverage: - functions_being_tested.append(dependent_func_coverage.name) - - graph = CoverageData._build_graph(main_func_coverage, dependent_func_coverage) - temp_json_file.unlink() - - return CoverageData( - file_path=source_code_path, - coverage=coverage, - function_name=function_name, - functions_being_tested=functions_being_tested, - graph=graph, - code_context=code_context, - main_func_coverage=main_func_coverage, - dependent_func_coverage=dependent_func_coverage, - status=status, - ) - - @staticmethod - def _parse_coverage_file( - coverage_file_path: Path, source_code_path: Path - ) -> tuple[dict[str, dict[str, Any]], CoverageStatus]: - with coverage_file_path.open() as f: - coverage_data = json.load(f) - - candidates = generate_candidates(source_code_path) - - logger.debug(f"Looking for coverage data in {' -> '.join(candidates)}") - for candidate in candidates: - try: - cov: dict[str, dict[str, Any]] = coverage_data["files"][candidate]["functions"] - logger.debug(f"Coverage data found for {source_code_path} in {candidate}") - status = CoverageStatus.PARSED_SUCCESSFULLY - break - except KeyError: - continue - else: - logger.debug(f"No coverage data found for {source_code_path} in {candidates}") - cov = {} - status = CoverageStatus.NOT_FOUND - return cov, status - - @staticmethod - def _fetch_function_coverages( - function_name: str, - code_context: CodeOptimizationContext, - coverage_data: dict[str, dict[str, Any]], - original_cov_data: dict[str, dict[str, Any]], - ) -> tuple[FunctionCoverage, Union[FunctionCoverage, None]]: - resolved_name = build_fully_qualified_name(function_name, code_context) - try: - main_function_coverage = FunctionCoverage( - name=resolved_name, - coverage=coverage_data[resolved_name]["summary"]["percent_covered"], - executed_lines=coverage_data[resolved_name]["executed_lines"], - unexecuted_lines=coverage_data[resolved_name]["missing_lines"], - executed_branches=coverage_data[resolved_name]["executed_branches"], - unexecuted_branches=coverage_data[resolved_name]["missing_branches"], - ) - except KeyError: - main_function_coverage = FunctionCoverage( - name=resolved_name, - coverage=0, - executed_lines=[], - unexecuted_lines=[], - executed_branches=[], - unexecuted_branches=[], - ) - - dependent_function = extract_dependent_function(function_name, code_context) - dependent_func_coverage = ( - CoverageData.grab_dependent_function_from_coverage_data( - dependent_function, coverage_data, original_cov_data - ) - if dependent_function - else None - ) - - return main_function_coverage, dependent_func_coverage - - @staticmethod - def _aggregate_coverage( - main_func_coverage: FunctionCoverage, dependent_func_coverage: Union[FunctionCoverage, None] - ) -> tuple[set[int], set[int]]: - total_executed_lines = set(main_func_coverage.executed_lines) - total_unexecuted_lines = set(main_func_coverage.unexecuted_lines) - - if dependent_func_coverage: - total_executed_lines.update(dependent_func_coverage.executed_lines) - total_unexecuted_lines.update(dependent_func_coverage.unexecuted_lines) - - return total_executed_lines, total_unexecuted_lines - - @staticmethod - def _build_graph( - main_func_coverage: FunctionCoverage, dependent_func_coverage: Union[FunctionCoverage, None] - ) -> dict[str, dict[str, Collection[object]]]: - graph = { - main_func_coverage.name: { - "executed_lines": set(main_func_coverage.executed_lines), - "unexecuted_lines": set(main_func_coverage.unexecuted_lines), - "executed_branches": main_func_coverage.executed_branches, - "unexecuted_branches": main_func_coverage.unexecuted_branches, - } - } - - if dependent_func_coverage: - graph[dependent_func_coverage.name] = { - "executed_lines": set(dependent_func_coverage.executed_lines), - "unexecuted_lines": set(dependent_func_coverage.unexecuted_lines), - "executed_branches": dependent_func_coverage.executed_branches, - "unexecuted_branches": dependent_func_coverage.unexecuted_branches, - } - - return graph - - @staticmethod - def grab_dependent_function_from_coverage_data( - dependent_function_name: str, - coverage_data: dict[str, dict[str, Any]], - original_cov_data: dict[str, dict[str, Any]], - ) -> FunctionCoverage: - """Grab the dependent function from the coverage data.""" - try: - return FunctionCoverage( - name=dependent_function_name, - coverage=coverage_data[dependent_function_name]["summary"]["percent_covered"], - executed_lines=coverage_data[dependent_function_name]["executed_lines"], - unexecuted_lines=coverage_data[dependent_function_name]["missing_lines"], - executed_branches=coverage_data[dependent_function_name]["executed_branches"], - unexecuted_branches=coverage_data[dependent_function_name]["missing_branches"], - ) - except KeyError: - msg = f"Coverage data not found for dependent function {dependent_function_name} in the coverage data" - try: - files = original_cov_data["files"] - for file in files: - functions = files[file]["functions"] - for function in functions: - if dependent_function_name in function: - return FunctionCoverage( - name=dependent_function_name, - coverage=functions[function]["summary"]["percent_covered"], - executed_lines=functions[function]["executed_lines"], - unexecuted_lines=functions[function]["missing_lines"], - executed_branches=functions[function]["executed_branches"], - unexecuted_branches=functions[function]["missing_branches"], - ) - msg = f"Coverage data not found for dependent function {dependent_function_name} in the original coverage data" - except KeyError: - raise ValueError(msg) from None - - return FunctionCoverage( - name=dependent_function_name, - coverage=0, - executed_lines=[], - unexecuted_lines=[], - executed_branches=[], - unexecuted_branches=[], - ) - def build_message(self) -> str: if self.status == CoverageStatus.NOT_FOUND: return f"No coverage data found for {self.function_name}" @@ -549,7 +346,6 @@ def create_empty(cls, file_path: Path, function_name: str, code_context: CodeOpt status=CoverageStatus.NOT_FOUND, ) - @dataclass class FunctionCoverage: """Represents the coverage data for a specific function in a source file.""" @@ -565,3 +361,249 @@ class FunctionCoverage: class TestingMode(enum.Enum): BEHAVIOR = "behavior" PERFORMANCE = "performance" + + +class VerificationType(str, Enum): + FUNCTION_CALL = ( + "function_call" # Correctness verification for a test function, checks input values and output values) + ) + INIT_STATE_FTO = "init_state_fto" # Correctness verification for fto class instance attributes after init + INIT_STATE_HELPER = "init_state_helper" # Correctness verification for helper class instance attributes after init + + +class TestType(Enum): + EXISTING_UNIT_TEST = 1 + INSPIRED_REGRESSION = 2 + GENERATED_REGRESSION = 3 + REPLAY_TEST = 4 + CONCOLIC_COVERAGE_TEST = 5 + INIT_STATE_TEST = 6 + + def to_name(self) -> str: + if self is TestType.INIT_STATE_TEST: + return "" + names = { + TestType.EXISTING_UNIT_TEST: "βš™οΈ Existing Unit Tests", + TestType.INSPIRED_REGRESSION: "🎨 Inspired Regression Tests", + TestType.GENERATED_REGRESSION: "πŸŒ€ Generated Regression Tests", + TestType.REPLAY_TEST: "βͺ Replay Tests", + TestType.CONCOLIC_COVERAGE_TEST: "πŸ”Ž Concolic Coverage Tests", + } + return names[self] + + +@dataclass(frozen=True) +class InvocationId: + test_module_path: str # The fully qualified name of the test module + test_class_name: Optional[str] # The name of the class where the test is defined + test_function_name: Optional[str] # The name of the test_function. Does not include the components of the file_name + function_getting_tested: str + iteration_id: Optional[str] + + # test_module_path:TestSuiteClass.test_function_name:function_tested:iteration_id + def id(self) -> str: + class_prefix = f"{self.test_class_name}." if self.test_class_name else "" + return ( + f"{self.test_module_path}:{class_prefix}{self.test_function_name}:" + f"{self.function_getting_tested}:{self.iteration_id}" + ) + + @staticmethod + def from_str_id(string_id: str, iteration_id: Optional[str] = None) -> InvocationId: + components = string_id.split(":") + assert len(components) == 4 + second_components = components[1].split(".") + if len(second_components) == 1: + test_class_name = None + test_function_name = second_components[0] + else: + test_class_name = second_components[0] + test_function_name = second_components[1] + return InvocationId( + test_module_path=components[0], + test_class_name=test_class_name, + test_function_name=test_function_name, + function_getting_tested=components[2], + iteration_id=iteration_id if iteration_id else components[3], + ) + + +@dataclass(frozen=True) +class FunctionTestInvocation: + loop_index: int # The loop index of the function invocation, starts at 1 + id: InvocationId # The fully qualified name of the function invocation (id) + file_name: Path # The file where the test is defined + did_pass: bool # Whether the test this function invocation was part of, passed or failed + runtime: Optional[int] # Time in nanoseconds + test_framework: str # unittest or pytest + test_type: TestType + return_value: Optional[object] # The return value of the function invocation + timed_out: Optional[bool] + verification_type: Optional[str] = VerificationType.FUNCTION_CALL + stdout: Optional[str] = None + + @property + def unique_invocation_loop_id(self) -> str: + return f"{self.loop_index}:{self.id.id()}" + + +class TestResults(BaseModel): + # don't modify these directly, use the add method + # also we don't support deletion of test results elements - caution is advised + test_results: list[FunctionTestInvocation] = [] + test_result_idx: dict[str, int] = {} + + def add(self, function_test_invocation: FunctionTestInvocation) -> None: + unique_id = function_test_invocation.unique_invocation_loop_id + if unique_id in self.test_result_idx: + if DEBUG_MODE: + logger.warning(f"Test result with id {unique_id} already exists. SKIPPING") + return + self.test_result_idx[unique_id] = len(self.test_results) + self.test_results.append(function_test_invocation) + + def merge(self, other: TestResults) -> None: + original_len = len(self.test_results) + self.test_results.extend(other.test_results) + for k, v in other.test_result_idx.items(): + if k in self.test_result_idx: + msg = f"Test result with id {k} already exists." + raise ValueError(msg) + self.test_result_idx[k] = v + original_len + + def group_by_benchmarks(self, benchmark_keys:list[BenchmarkKey], benchmark_replay_test_dir: Path, project_root: Path) -> dict[BenchmarkKey, TestResults]: + """Group TestResults by benchmark for calculating improvements for each benchmark.""" + + 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:]}_{benchmark_key.function_name}__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(): + if test_result.id.test_module_path.startswith(module_path): + test_results_by_benchmark[benchmark_key].add(test_result) + + return test_results_by_benchmark + + def get_by_unique_invocation_loop_id(self, unique_invocation_loop_id: str) -> FunctionTestInvocation | None: + try: + return self.test_results[self.test_result_idx[unique_invocation_loop_id]] + except (IndexError, KeyError): + return None + + def get_all_ids(self) -> set[InvocationId]: + return {test_result.id for test_result in self.test_results} + + def get_all_unique_invocation_loop_ids(self) -> set[str]: + return {test_result.unique_invocation_loop_id for test_result in self.test_results} + + def number_of_loops(self) -> int: + if not self.test_results: + return 0 + return max(test_result.loop_index for test_result in self.test_results) + + def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]: + report = {} + for test_type in TestType: + report[test_type] = {"passed": 0, "failed": 0} + for test_result in self.test_results: + if test_result.loop_index == 1: + if test_result.did_pass: + report[test_result.test_type]["passed"] += 1 + else: + report[test_result.test_type]["failed"] += 1 + return report + + @staticmethod + def report_to_string(report: dict[TestType, dict[str, int]]) -> str: + return " ".join( + [ + f"{test_type.to_name()}- (Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']})" + for test_type in TestType + ] + ) + + @staticmethod + def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree: + tree = Tree(title) + for test_type in TestType: + if test_type is TestType.INIT_STATE_TEST: + continue + tree.add( + f"{test_type.to_name()} - Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']}" + ) + return + + def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: + + usable_runtime = defaultdict(list) + for result in self.test_results: + if result.did_pass: + if not result.runtime: + msg = ( + f"Ignoring test case that passed but had no runtime -> {result.id}, " + f"Loop # {result.loop_index}, Test Type: {result.test_type}, " + f"Verification Type: {result.verification_type}" + ) + logger.debug(msg) + else: + usable_runtime[result.id].append(result.runtime) + return usable_runtime + + def total_passed_runtime(self) -> int: + """Calculate the sum of runtimes of all test cases that passed. + + A testcase runtime is the minimum value of all looped execution runtimes. + + :return: The runtime in nanoseconds. + """ + return sum( + [min(usable_runtime_data) for _, usable_runtime_data in self.usable_runtime_data_by_test_case().items()] + ) + + def __iter__(self) -> Iterator[FunctionTestInvocation]: + return iter(self.test_results) + + def __len__(self) -> int: + return len(self.test_results) + + def __getitem__(self, index: int) -> FunctionTestInvocation: + return self.test_results[index] + + def __setitem__(self, index: int, value: FunctionTestInvocation) -> None: + self.test_results[index] = value + + def __contains__(self, value: FunctionTestInvocation) -> bool: + return value in self.test_results + + def __bool__(self) -> bool: + return bool(self.test_results) + + def __eq__(self, other: object) -> bool: + # Unordered comparison + if type(self) is not type(other): + return False + if len(self) != len(other): + return False + original_recursion_limit = sys.getrecursionlimit() + cast(TestResults, other) + for test_result in self: + other_test_result = other.get_by_unique_invocation_loop_id(test_result.unique_invocation_loop_id) + if other_test_result is None: + return False + + if original_recursion_limit < 5000: + sys.setrecursionlimit(5000) + if ( + test_result.file_name != other_test_result.file_name + or test_result.did_pass != other_test_result.did_pass + or test_result.runtime != other_test_result.runtime + or test_result.test_framework != other_test_result.test_framework + or test_result.test_type != other_test_result.test_type + or not comparator(test_result.return_value, other_test_result.return_value) + ): + sys.setrecursionlimit(original_recursion_limit) + return False + sys.setrecursionlimit(original_recursion_limit) + return True \ No newline at end of file diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index c0143eb10..807fd3a8c 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -22,12 +22,12 @@ from codeflash.benchmarking.utils import process_benchmark_data from codeflash.cli_cmds.console import code_print, console, logger, progress_bar from codeflash.code_utils import env_utils -from codeflash.code_utils.code_extractor import add_needed_imports_from_module, extract_code from codeflash.code_utils.code_replacer import replace_function_definitions_in_module from codeflash.code_utils.code_utils import ( cleanup_paths, file_name_from_test_module_name, get_run_tmp_file, + has_any_async_functions, module_name_from_file_path, ) from codeflash.code_utils.config_consts import ( @@ -49,7 +49,6 @@ BestOptimization, CodeOptimizationContext, FunctionCalledInTest, - FunctionParent, GeneratedTests, GeneratedTestsList, OptimizationSet, @@ -58,8 +57,9 @@ TestFile, TestFiles, TestingMode, + TestResults, + TestType, BenchmarkKey, ) -from codeflash.optimization.function_context import get_constrained_function_context_and_helper_functions from codeflash.result.create_pr import check_create_pr, existing_tests_source_for from codeflash.result.critic import coverage_critic, performance_gain, quantity_of_tests_critic, speedup_critic from codeflash.result.explanation import Explanation @@ -68,7 +68,6 @@ from codeflash.verification.equivalence import compare_test_results from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture from codeflash.verification.parse_test_output import parse_test_results -from codeflash.verification.test_results import TestResults, TestType from codeflash.verification.test_runner import run_behavioral_tests, run_benchmarking_tests from codeflash.verification.verification_utils import get_test_file_path from codeflash.verification.verifier import generate_tests @@ -76,9 +75,6 @@ if TYPE_CHECKING: from argparse import Namespace - import numpy as np - import numpy.typing as npt - from codeflash.either import Result from codeflash.models.models import CoverageData, FunctionSource, OptimizedCandidate from codeflash.verification.verification_utils import TestConfig @@ -163,7 +159,7 @@ def optimize_function(self) -> Result[BestOptimization, str]: transient=True, ): generated_results = self.generate_tests_and_optimizations( - code_to_optimize_with_helpers=code_context.code_to_optimize_with_helpers, + testgen_context_code=code_context.testgen_context_code, read_writable_code=code_context.read_writable_code, read_only_context_code=code_context.read_only_context_code, helper_functions=code_context.helper_functions, @@ -242,7 +238,7 @@ def optimize_function(self) -> Result[BestOptimization, str]: best_optimization = None - for u, candidates in enumerate([optimizations_set.control, optimizations_set.experiment]): + for _u, candidates in enumerate([optimizations_set.control, optimizations_set.experiment]): if candidates is None: continue @@ -432,30 +428,30 @@ def determine_best_candidate( ) tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") + replay_perf_gain = {} if self.args.benchmark: - - benchmark_keys = {(benchmark.file_name, benchmark.function_name) for benchmark in self.total_benchmark_timings} - test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmark(benchmark_keys) - for benchmark_key, test_results in test_results_by_benchmark.items(): + logger.info(f"Calculating benchmark improvement..") + test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) + for benchmark_key, candidate_test_results in test_results_by_benchmark.items(): original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results[benchmark_key].total_passed_runtime() - candidate_replay_runtime = candidate_result.replay_benchmarking_test_results.total_passed_runtime() - replay_perf_gain = performance_gain( + candidate_replay_runtime = candidate_test_results.total_passed_runtime() + replay_perf_gain[benchmark_key] = performance_gain( original_runtime_ns=original_code_replay_runtime, optimized_runtime_ns=candidate_replay_runtime, ) - tree.add(f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") + tree.add( + f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") tree.add( f"Best benchmark replay runtime: {humanize_runtime(candidate_replay_runtime)} " f"(measured over {candidate_result.max_loop_count} " f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" ) - tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain + 1:.1f}X") + tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain[benchmark_key] * 100:.1f}%") + tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain[benchmark_key] + 1:.1f}X") best_optimization = BestOptimization( candidate=candidate, helper_functions=code_context.helper_functions, runtime=best_test_runtime, - replay_runtime=candidate_replay_runtime if self.args.benchmark else None, winning_behavioral_test_results=candidate_result.behavior_test_results, replay_performance_gain=replay_perf_gain if self.args.benchmark else None, winning_benchmarking_test_results=candidate_result.benchmarking_test_results, @@ -582,50 +578,6 @@ def replace_function_and_helpers_with_optimized_code( return did_update def get_code_optimization_context(self) -> Result[CodeOptimizationContext, str]: - code_to_optimize, contextual_dunder_methods = extract_code([self.function_to_optimize]) - if code_to_optimize is None: - return Failure("Could not find function to optimize.") - (helper_code, helper_functions, helper_dunder_methods) = get_constrained_function_context_and_helper_functions( - self.function_to_optimize, self.project_root, code_to_optimize - ) - if self.function_to_optimize.parents: - function_class = self.function_to_optimize.parents[0].name - same_class_helper_methods = [ - df - for df in helper_functions - if df.qualified_name.count(".") > 0 and df.qualified_name.split(".")[0] == function_class - ] - optimizable_methods = [ - FunctionToOptimize( - df.qualified_name.split(".")[-1], - df.file_path, - [FunctionParent(df.qualified_name.split(".")[0], "ClassDef")], - None, - None, - ) - for df in same_class_helper_methods - ] + [self.function_to_optimize] - dedup_optimizable_methods = [] - added_methods = set() - for method in reversed(optimizable_methods): - if f"{method.file_path}.{method.qualified_name}" not in added_methods: - dedup_optimizable_methods.append(method) - added_methods.add(f"{method.file_path}.{method.qualified_name}") - if len(dedup_optimizable_methods) > 1: - code_to_optimize, contextual_dunder_methods = extract_code(list(reversed(dedup_optimizable_methods))) - if code_to_optimize is None: - return Failure("Could not find function to optimize.") - code_to_optimize_with_helpers = helper_code + "\n" + code_to_optimize - - code_to_optimize_with_helpers_and_imports = add_needed_imports_from_module( - self.function_to_optimize_source_code, - code_to_optimize_with_helpers, - self.function_to_optimize.file_path, - self.function_to_optimize.file_path, - self.project_root, - helper_functions, - ) - try: new_code_ctx = code_context_extractor.get_code_optimization_context( self.function_to_optimize, self.project_root @@ -635,7 +587,7 @@ def get_code_optimization_context(self) -> Result[CodeOptimizationContext, str]: return Success( CodeOptimizationContext( - code_to_optimize_with_helpers=code_to_optimize_with_helpers_and_imports, + testgen_context_code=new_code_ctx.testgen_context_code, read_writable_code=new_code_ctx.read_writable_code, read_only_context_code=new_code_ctx.read_only_context_code, helper_functions=new_code_ctx.helper_functions, # only functions that are read writable @@ -739,7 +691,7 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, list[Functi def generate_tests_and_optimizations( self, - code_to_optimize_with_helpers: str, + testgen_context_code: str, read_writable_code: str, read_only_context_code: str, helper_functions: list[FunctionSource], @@ -754,7 +706,7 @@ def generate_tests_and_optimizations( # Submit the test generation task as future future_tests = self.generate_and_instrument_tests( executor, - code_to_optimize_with_helpers, + testgen_context_code, [definition.fully_qualified_name for definition in helper_functions], generated_test_paths, generated_perf_test_paths, @@ -883,9 +835,7 @@ def establish_original_code_baseline( ) console.rule() return Failure("Failed to establish a baseline for the original code - bevhavioral tests failed.") - if not coverage_critic( - coverage_results, self.args.test_framework - ): + if not coverage_critic(coverage_results, self.args.test_framework): return Failure("The threshold for test coverage was not met.") if test_framework == "pytest": benchmarking_results, _ = self.run_and_parse_tests( @@ -953,8 +903,10 @@ def establish_original_code_baseline( logger.debug(f"Total original code runtime (ns): {total_timing}") if self.args.benchmark: - replay_benchmarking_test_results = benchmarking_results.filter_by_test_type(TestType.REPLAY_TEST) - logger.info(f"Total replay test runtime: {humanize_runtime(replay_benchmarking_test_results.total_passed_runtime())}") + replay_benchmarking_test_results = benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) + for benchmark_name, benchmark_results in replay_benchmarking_test_results.items(): + + logger.info(f"Replay benchmark '{benchmark_name}' runtime: {humanize_runtime(benchmark_results.total_passed_runtime())}") return Success( ( OriginalCodeBaseline( @@ -1075,10 +1027,9 @@ def run_optimized_candidate( logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") if self.args.benchmark: - candidate_replay_benchmarking_results = candidate_benchmarking_results.filter_by_test_type(TestType.REPLAY_TEST) - logger.debug( - f"Total optimized code {optimization_candidate_index} replay benchmark runtime (ns): {candidate_replay_benchmarking_results.total_passed_runtime()}" - ) + candidate_replay_benchmarking_results = candidate_benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) + for benchmark_name, benchmark_results in candidate_replay_benchmarking_results.items(): + logger.debug(f"Benchmark {benchmark_name} runtime (ns): {humanize_runtime(benchmark_results.total_passed_runtime())}") return Success( OptimizedCandidateResult( max_loop_count=loop_count, @@ -1106,15 +1057,15 @@ def run_and_parse_tests( unittest_loop_index: int | None = None, ) -> tuple[TestResults, CoverageData | None]: coverage_database_file = None + coverage_config_file = None try: if testing_type == TestingMode.BEHAVIOR: - result_file_path, run_result, coverage_database_file = run_behavioral_tests( + result_file_path, run_result, coverage_database_file, coverage_config_file = run_behavioral_tests( test_files, test_framework=self.test_cfg.test_framework, cwd=self.project_root, test_env=test_env, pytest_timeout=INDIVIDUAL_TESTCASE_TIMEOUT, - pytest_cmd=self.test_cfg.pytest_cmd, verbose=True, enable_coverage=enable_coverage, ) @@ -1123,24 +1074,25 @@ def run_and_parse_tests( test_files, cwd=self.project_root, test_env=test_env, - pytest_timeout=INDIVIDUAL_TESTCASE_TIMEOUT, pytest_cmd=self.test_cfg.pytest_cmd, + pytest_timeout=INDIVIDUAL_TESTCASE_TIMEOUT, pytest_target_runtime_seconds=testing_time, pytest_min_loops=pytest_min_loops, pytest_max_loops=pytest_max_loops, test_framework=self.test_cfg.test_framework, ) else: - raise ValueError(f"Unexpected testing type: {testing_type}") + msg = f"Unexpected testing type: {testing_type}" + raise ValueError(msg) except subprocess.TimeoutExpired: logger.exception( - f'Error running tests in {", ".join(str(f) for f in test_files.test_files)}.\nTimeout Error' + f"Error running tests in {', '.join(str(f) for f in test_files.test_files)}.\nTimeout Error" ) return TestResults(), None if run_result.returncode != 0 and testing_type == TestingMode.BEHAVIOR: logger.debug( - f'Nonzero return code {run_result.returncode} when running tests in ' - f'{", ".join([str(f.instrumented_behavior_file_path) for f in test_files.test_files])}.\n' + f"Nonzero return code {run_result.returncode} when running tests in " + f"{', '.join([str(f.instrumented_behavior_file_path) for f in test_files.test_files])}.\n" f"stdout: {run_result.stdout}\n" f"stderr: {run_result.stderr}\n" ) @@ -1155,6 +1107,7 @@ def run_and_parse_tests( source_file=self.function_to_optimize.file_path, code_context=code_context, coverage_database_file=coverage_database_file, + coverage_config_file=coverage_config_file, ) return results, coverage_results @@ -1185,4 +1138,3 @@ def generate_and_instrument_tests( zip(generated_test_paths, generated_perf_test_paths) ) ] - diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 35d91a274..51b2ea29e 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -48,6 +48,7 @@ def __init__(self, args: Namespace) -> None: project_root_path=args.project_root, test_framework=args.test_framework, pytest_cmd=args.pytest_cmd, + benchmark_tests_root=args.benchmarks_root if "benchmark" in args and "benchmarks_root" in args else None, ) self.aiservice_client = AiServiceClient() @@ -114,7 +115,7 @@ def run(self) -> None: if trace_file.exists(): trace_file.unlink() - replay_tests_dir = Path(self.args.tests_root) + replay_tests_dir = Path(self.args.benchmarks_root) / "codeflash_replay_tests" trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Run all tests that use pytest-benchmark replay_count = generate_replay_test(trace_file, replay_tests_dir) if replay_count == 0: @@ -251,8 +252,8 @@ def run(self) -> None: if function_optimizer.test_cfg.concolic_test_root_dir: shutil.rmtree(function_optimizer.test_cfg.concolic_test_root_dir, ignore_errors=True) if self.args.benchmark: - if replay_tests_dir.exists(): - shutil.rmtree(replay_tests_dir, ignore_errors=True) + # if replay_tests_dir.exists(): + # shutil.rmtree(replay_tests_dir, ignore_errors=True) trace_file.unlink(missing_ok=True) if hasattr(get_run_tmp_file, "tmpdir"): get_run_tmp_file.tmpdir.cleanup() diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 10794991a..e56558a94 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -5,8 +5,7 @@ from pydantic.dataclasses import dataclass from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.models.models import BenchmarkDetail -from codeflash.verification.test_results import TestResults +from codeflash.models.models import BenchmarkDetail, TestResults @dataclass(frozen=True, config={"arbitrary_types_allowed": True}) diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index f08b2485a..c67c7e87d 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -27,7 +27,7 @@ def test_trace_benchmarks(): # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_name, benchmark_function_name, function_name") + "SELECT function_name, class_name, module_name, file_path, benchmark_function_name, benchmark_file_path, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_path, benchmark_function_name, function_name") function_calls = cursor.fetchall() # Assert the length of function calls @@ -39,44 +39,44 @@ def test_trace_benchmarks(): expected_calls = [ ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", "test_benchmark_bubble_sort.py", 20), + "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 20), ("sort_class", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", "test_benchmark_bubble_sort.py", 18), + "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 18), ("sort_static", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", "test_benchmark_bubble_sort.py", 19), + "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 19), ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", "test_benchmark_bubble_sort.py", 17), + "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 17), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_sort", "test_benchmark_bubble_sort.py", 7), + "test_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 7), ("compute_and_sort", "", "code_to_optimize.process_and_bubble_sort_codeflash_trace", f"{process_and_bubble_sort_path}", - "test_compute_and_sort", "test_process_and_sort.py", 4), + "test_compute_and_sort", str(benchmarks_root / "test_process_and_sort.py"), 4), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_no_func", "test_process_and_sort.py", 8), + "test_no_func", str(benchmarks_root / "test_process_and_sort.py"), 8), ] 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_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_file_name" + assert actual[5] == expected[5], f"Mismatch at index {idx} for benchmark_file_path" assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number" # Close connection conn.close() generate_replay_test(output_file, tests_root) - test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort_py_test_class_sort__replay_test_0.py") + test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort_test_class_sort__replay_test_0.py") assert test_class_sort_path.exists() test_class_sort_code = f""" import dill as pickle @@ -89,7 +89,7 @@ def test_trace_benchmarks(): trace_file_path = r"{output_file.as_posix()}" def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "sorter" @@ -102,7 +102,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): ret = instance.sorter(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) if not args: @@ -110,13 +110,13 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_class(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_static(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_static(*args, **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_name=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "__init__" @@ -131,7 +131,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): """ assert test_class_sort_path.read_text("utf-8").strip()==test_class_sort_code.strip() - test_sort_path = tests_root / Path("test_benchmark_bubble_sort_py_test_sort__replay_test_0.py") + test_sort_path = tests_root / Path("test_benchmark_bubble_sort_test_sort__replay_test_0.py") assert test_sort_path.exists() test_sort_code = f""" import dill as pickle @@ -144,7 +144,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): trace_file_path = r"{output_file}" def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_name=r"{bubble_sort_path}", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_path=r"{bubble_sort_path}", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_sorter(*args, **kwargs) @@ -172,7 +172,7 @@ def test_trace_multithreaded_benchmark() -> None: # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_name, benchmark_function_name, benchmark_file_name, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_name, benchmark_function_name, function_name") + "SELECT function_name, class_name, module_name, file_path, benchmark_function_name, benchmark_file_path, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_path, benchmark_function_name, function_name") function_calls = cursor.fetchall() # Assert the length of function calls @@ -192,15 +192,15 @@ def test_trace_multithreaded_benchmark() -> None: expected_calls = [ ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_benchmark_sort", "test_multithread_sort.py", 4), + "test_benchmark_sort", str(benchmarks_root / "test_multithread_sort.py"), 4), ] 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_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_file_name" + assert actual[5] == expected[5], f"Mismatch at index {idx} for benchmark_file_path" assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number" # Close connection conn.close() diff --git a/tests/test_unit_test_discovery.py b/tests/test_unit_test_discovery.py index c05b79e63..8c3bc35c8 100644 --- a/tests/test_unit_test_discovery.py +++ b/tests/test_unit_test_discovery.py @@ -3,7 +3,6 @@ from pathlib import Path from codeflash.discovery.discover_unit_tests import discover_unit_tests -from codeflash.verification.test_results import TestType from codeflash.verification.verification_utils import TestConfig From 0937329d386b64f89a871b34a134d7b0430a8bb9 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Mon, 31 Mar 2025 20:20:57 -0700 Subject: [PATCH 086/122] PRAGMA journal to memory to make it faster --- codeflash/benchmarking/codeflash_trace.py | 1 + codeflash/benchmarking/plugin/plugin.py | 1 + 2 files changed, 2 insertions(+) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 8c307f8a0..bcbb3268c 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -31,6 +31,7 @@ def setup(self, trace_path: str) -> None: self._connection = sqlite3.connect(self._trace_path) cur = self._connection.cursor() cur.execute("PRAGMA synchronous = OFF") + cur.execute("PRAGMA journal_mode = MEMORY") cur.execute( "CREATE TABLE IF NOT EXISTS benchmark_function_timings(" "function_name TEXT, class_name TEXT, module_name TEXT, file_path TEXT," diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index 09858601c..b022f9afb 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -22,6 +22,7 @@ def setup(self, trace_path:str) -> None: self._connection = sqlite3.connect(self._trace_path) cur = self._connection.cursor() cur.execute("PRAGMA synchronous = OFF") + 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," From ed8f5efcad07b417eaf515c5c7bda060df29d111 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 1 Apr 2025 10:53:14 -0700 Subject: [PATCH 087/122] benchmarks root must be subdir of tests root --- .../tests/unittest/test_bubble_sort.py | 36 +++++++++---------- .../unittest/test_bubble_sort_parametrized.py | 36 +++++++++---------- codeflash/cli_cmds/cli.py | 27 +++++++++----- 3 files changed, 55 insertions(+), 44 deletions(-) diff --git a/code_to_optimize/tests/unittest/test_bubble_sort.py b/code_to_optimize/tests/unittest/test_bubble_sort.py index 4c76414ef..200f82b7a 100644 --- a/code_to_optimize/tests/unittest/test_bubble_sort.py +++ b/code_to_optimize/tests/unittest/test_bubble_sort.py @@ -1,18 +1,18 @@ -# import unittest -# -# from code_to_optimize.bubble_sort import sorter -# -# -# class TestPigLatin(unittest.TestCase): -# def test_sort(self): -# input = [5, 4, 3, 2, 1, 0] -# output = sorter(input) -# self.assertEqual(output, [0, 1, 2, 3, 4, 5]) -# -# input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0] -# output = sorter(input) -# self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) -# -# input = list(reversed(range(5000))) -# output = sorter(input) -# self.assertEqual(output, list(range(5000))) +import unittest + +from code_to_optimize.bubble_sort import sorter + + +class TestPigLatin(unittest.TestCase): + def test_sort(self): + input = [5, 4, 3, 2, 1, 0] + output = sorter(input) + self.assertEqual(output, [0, 1, 2, 3, 4, 5]) + + input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0] + output = sorter(input) + self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) + + input = list(reversed(range(5000))) + output = sorter(input) + self.assertEqual(output, list(range(5000))) diff --git a/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py b/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py index c1aef993b..59c86abc8 100644 --- a/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py +++ b/code_to_optimize/tests/unittest/test_bubble_sort_parametrized.py @@ -1,18 +1,18 @@ -# import unittest -# -# from parameterized import parameterized -# -# from code_to_optimize.bubble_sort import sorter -# -# -# class TestPigLatin(unittest.TestCase): -# @parameterized.expand( -# [ -# ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]), -# ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), -# (list(reversed(range(50))), list(range(50))), -# ] -# ) -# def test_sort(self, input, expected_output): -# output = sorter(input) -# self.assertEqual(output, expected_output) +import unittest + +from parameterized import parameterized + +from code_to_optimize.bubble_sort import sorter + + +class TestPigLatin(unittest.TestCase): + @parameterized.expand( + [ + ([5, 4, 3, 2, 1, 0], [0, 1, 2, 3, 4, 5]), + ([5.0, 4.0, 3.0, 2.0, 1.0, 0.0], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]), + (list(reversed(range(50))), list(range(50))), + ] + ) + def test_sort(self, input, expected_output): + output = sorter(input) + self.assertEqual(output, expected_output) diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index d1e786703..07652f707 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -135,14 +135,25 @@ def process_pyproject_config(args: Namespace) -> Namespace: if args.benchmark: assert args.benchmarks_root is not None, "--benchmarks-root must be specified when running with --benchmark" assert Path(args.benchmarks_root).is_dir(), f"--benchmarks-root {args.benchmarks_root} must be a valid directory" - assert not (env_utils.get_pr_number() is not None and not env_utils.ensure_codeflash_api_key()), ( - "Codeflash API key not found. When running in a Github Actions Context, provide the " - "'CODEFLASH_API_KEY' environment variable as a secret.\n" - "You can add a secret by going to your repository's settings page, then clicking 'Secrets' in the left sidebar.\n" - "Then, click 'New repository secret' and add your api key with the variable name CODEFLASH_API_KEY.\n" - f"Here's a direct link: {get_github_secrets_page_url()}\n" - "Exiting..." - ) + assert Path(args.benchmarks_root).is_relative_to(Path(args.tests_root)), ( + f"--benchmarks-root {args.benchmarks_root} must be a subdirectory of --tests-root {args.tests_root}" + ) + if env_utils.get_pr_number() is not None: + assert env_utils.ensure_codeflash_api_key(), ( + "Codeflash API key not found. When running in a Github Actions Context, provide the " + "'CODEFLASH_API_KEY' environment variable as a secret.\n" + "You can add a secret by going to your repository's settings page, then clicking 'Secrets' in the left sidebar.\n" + "Then, click 'New repository secret' and add your api key with the variable name CODEFLASH_API_KEY.\n" + f"Here's a direct link: {get_github_secrets_page_url()}\n" + "Exiting..." + ) + + repo = git.Repo(search_parent_directories=True) + + owner, repo_name = get_repo_owner_and_name(repo) + + require_github_app_or_exit(owner, repo_name) + if hasattr(args, "ignore_paths") and args.ignore_paths is not None: normalized_ignore_paths = [] for path in args.ignore_paths: From 75c1be7b1102d1746646903119548853885743f3 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 1 Apr 2025 14:44:17 -0700 Subject: [PATCH 088/122] replay tests are now grouped by benchmark file. each benchmark test file will create one replay test file. --- .../test_benchmark_bubble_sort.py | 6 ++ codeflash/benchmarking/replay_test.py | 89 ++++++++++--------- codeflash/models/models.py | 2 +- tests/test_trace_benchmarks.py | 47 ++++++---- 4 files changed, 86 insertions(+), 58 deletions(-) diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py index 03b9d38d1..21f9755a5 100644 --- a/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py +++ b/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py @@ -15,6 +15,12 @@ def test_sort2(): def test_class_sort(benchmark): obj = Sorter(list(reversed(range(100)))) result1 = benchmark(obj.sorter, 2) + +def test_class_sort2(benchmark): result2 = benchmark(Sorter.sort_class, list(reversed(range(100)))) + +def test_class_sort3(benchmark): result3 = benchmark(Sorter.sort_static, list(reversed(range(100)))) + +def test_class_sort4(benchmark): result4 = benchmark(Sorter, [1,2,3]) \ No newline at end of file diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 9ecac2ec4..6466b24db 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -16,7 +16,7 @@ def get_next_arg_and_return( - trace_file: str, function_name: str, file_path: str, class_name: str | None = None, num_to_get: int = 256 + trace_file: str, benchmark_function_name:str, function_name: str, file_path: str, class_name: str | None = None, num_to_get: int = 256 ) -> Generator[Any]: db = sqlite3.connect(trace_file) cur = db.cursor() @@ -24,13 +24,13 @@ def get_next_arg_and_return( if class_name is not None: cursor = cur.execute( - "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_path = ? AND class_name = ? LIMIT ?", - (function_name, file_path, class_name, limit), + "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = ? LIMIT ?", + (benchmark_function_name, function_name, file_path, class_name, limit), ) else: cursor = cur.execute( - "SELECT * FROM benchmark_function_timings WHERE function_name = ? AND file_path = ? AND class_name = '' LIMIT ?", - (function_name, file_path, limit), + "SELECT * FROM benchmark_function_timings WHERE benchmark_function_name = ? AND function_name = ? AND file_path = ? AND class_name = '' LIMIT ?", + (benchmark_function_name, function_name, file_path, limit), ) while (val := cursor.fetchone()) is not None: @@ -61,6 +61,7 @@ def create_trace_replay_test_code( """ assert test_framework in ["pytest", "unittest"] + # Create Imports imports = f"""import dill as pickle {"import unittest" if test_framework == "unittest" else ""} from codeflash.benchmarking.replay_test import get_next_arg_and_return @@ -82,16 +83,15 @@ def create_trace_replay_test_code( imports += "\n".join(function_imports) - functions_to_optimize = [func.get("function_name") for func in functions_data - if func.get("function_name") != "__init__"] + functions_to_optimize = sorted({func.get("function_name") for func in functions_data + if func.get("function_name") != "__init__"}) metadata = f"""functions = {functions_to_optimize} trace_file_path = r"{trace_file}" """ - # Templates for different types of tests test_function_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = {function_name}(*args, **kwargs) @@ -100,7 +100,7 @@ def create_trace_replay_test_code( test_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} function_name = "{orig_function_name}" @@ -115,7 +115,7 @@ def create_trace_replay_test_code( test_class_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} if not args: @@ -125,13 +125,15 @@ def create_trace_replay_test_code( ) test_static_method_body = textwrap.dedent( """\ - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="{benchmark_function_name}", function_name="{orig_function_name}", file_path=r"{file_path}", class_name="{class_name}", num_to_get={max_run_count}): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl){filter_variables} ret = {class_name_alias}{method_name}(*args, **kwargs) """ ) + # Create main body + if test_framework == "unittest": self = "self" test_template = "\nclass TestTracedFunctions(unittest.TestCase):\n" @@ -140,17 +142,20 @@ def create_trace_replay_test_code( self = "" for func in functions_data: + module_name = func.get("module_name") function_name = func.get("function_name") class_name = func.get("class_name") file_path = func.get("file_path") + benchmark_function_name = func.get("benchmark_function_name") function_properties = func.get("function_properties") if not class_name: alias = get_function_alias(module_name, function_name) test_body = test_function_body.format( + benchmark_function_name=benchmark_function_name, + orig_function_name=function_name, function_name=alias, file_path=file_path, - orig_function_name=function_name, max_run_count=max_run_count, ) else: @@ -162,6 +167,7 @@ def create_trace_replay_test_code( method_name = "." + function_name if function_name != "__init__" else "" if function_properties.is_classmethod: test_body = test_class_method_body.format( + benchmark_function_name=benchmark_function_name, orig_function_name=function_name, file_path=file_path, class_name_alias=class_name_alias, @@ -172,6 +178,7 @@ def create_trace_replay_test_code( ) elif function_properties.is_staticmethod: test_body = test_static_method_body.format( + benchmark_function_name=benchmark_function_name, orig_function_name=function_name, file_path=file_path, class_name_alias=class_name_alias, @@ -182,6 +189,7 @@ def create_trace_replay_test_code( ) else: test_body = test_method_body.format( + benchmark_function_name=benchmark_function_name, orig_function_name=function_name, file_path=file_path, class_name_alias=class_name_alias, @@ -217,25 +225,25 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework conn = sqlite3.connect(trace_file_path.as_posix()) cursor = conn.cursor() - # Get distinct benchmark names + # Get distinct benchmark file paths cursor.execute( - "SELECT DISTINCT benchmark_function_name, benchmark_file_path FROM benchmark_function_timings" + "SELECT DISTINCT benchmark_file_path FROM benchmark_function_timings" ) - benchmarks = cursor.fetchall() + benchmark_files = cursor.fetchall() - # Generate a test for each benchmark - for benchmark in benchmarks: - benchmark_function_name, benchmark_file_path = benchmark - # Get functions associated with this benchmark + # Generate a test for each benchmark file + for benchmark_file in benchmark_files: + benchmark_file_path = benchmark_file[0] + # Get all benchmarks and functions associated with this file path cursor.execute( - "SELECT DISTINCT function_name, class_name, module_name, file_path, benchmark_line_number FROM benchmark_function_timings " - "WHERE benchmark_function_name = ? AND benchmark_file_path = ?", - (benchmark_function_name, benchmark_file_path) + "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,) ) functions_data = [] - for func_row in cursor.fetchall(): - function_name, class_name, module_name, file_path, benchmark_line_number = func_row + for row in cursor.fetchall(): + benchmark_function_name, function_name, class_name, module_name, file_path, benchmark_line_number = row # Add this function to our list functions_data.append({ "function_name": function_name, @@ -246,16 +254,15 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework "benchmark_file_path": benchmark_file_path, "benchmark_line_number": benchmark_line_number, "function_properties": inspect_top_level_functions_or_methods( - file_name=Path(file_path), - function_or_method_name=function_name, - class_name=class_name, - ) + file_name=Path(file_path), + function_or_method_name=function_name, + class_name=class_name, + ) }) if not functions_data: - logger.info(f"No functions found for benchmark {benchmark_function_name} in {benchmark_file_path}") + logger.info(f"No benchmark test functions found in {benchmark_file_path}") continue - # Generate the test code for this benchmark test_code = create_trace_replay_test_code( trace_file=trace_file_path.as_posix(), @@ -265,17 +272,15 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework ) test_code = isort.code(test_code) - # Write to file if requested - if output_dir: - name = Path(benchmark_file_path).name.split(".")[0][5:] # remove "test_" from the name since we add it in later - output_file = get_test_file_path( - test_dir=Path(output_dir), function_name=f"{name}_{benchmark_function_name}", test_type="replay" - ) - # Write test code to file, parents = true - output_dir.mkdir(parents=True, exist_ok=True) - output_file.write_text(test_code, "utf-8") - count += 1 - logger.info(f"Replay test for benchmark `{benchmark_function_name}` in {name} written to {output_file}") + name = Path(benchmark_file_path).name.split(".")[0][5:] # remove "test_" from the name since we add it in later + output_file = get_test_file_path( + test_dir=Path(output_dir), function_name=f"{name}", test_type="replay" + ) + # Write test code to file, parents = true + output_dir.mkdir(parents=True, exist_ok=True) + output_file.write_text(test_code, "utf-8") + count += 1 + logger.info(f"Replay test for benchmark file `{benchmark_file_path}` in {name} written to {output_file}") conn.close() diff --git a/codeflash/models/models.py b/codeflash/models/models.py index ed0360eef..e415108e7 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -477,7 +477,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:]}_{benchmark_key.function_name}__replay_test_", project_root) + 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) 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(): diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index c67c7e87d..7851230d9 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -37,21 +37,21 @@ def test_trace_benchmarks(): process_and_bubble_sort_path = (project_root / "process_and_bubble_sort_codeflash_trace.py").as_posix() # Expected function calls expected_calls = [ - ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 20), + "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 17), ("sort_class", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 18), + "test_class_sort2", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 20), ("sort_static", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 19), + "test_class_sort3", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 23), - ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", + ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 17), + "test_class_sort4", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 26), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", @@ -76,20 +76,28 @@ def test_trace_benchmarks(): # Close connection conn.close() generate_replay_test(output_file, tests_root) - test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort_test_class_sort__replay_test_0.py") + test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort__replay_test_0.py") assert test_class_sort_path.exists() test_class_sort_code = f""" import dill as pickle from code_to_optimize.bubble_sort_codeflash_trace import \\ Sorter as code_to_optimize_bubble_sort_codeflash_trace_Sorter +from code_to_optimize.bubble_sort_codeflash_trace import \\ + sorter as code_to_optimize_bubble_sort_codeflash_trace_sorter from codeflash.benchmarking.replay_test import get_next_arg_and_return -functions = ['sorter', 'sort_class', 'sort_static'] +functions = ['sort_class', 'sort_static', 'sorter'] trace_file_path = r"{output_file.as_posix()}" +def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_sort", function_name="sorter", file_path=r"{bubble_sort_path}", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + ret = code_to_optimize_bubble_sort_codeflash_trace_sorter(*args, **kwargs) + def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_class_sort", function_name="sorter", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "sorter" @@ -102,7 +110,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sorter(): ret = instance.sorter(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_class", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_class_sort2", function_name="sort_class", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) if not args: @@ -110,13 +118,13 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_class(): ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_class(*args[1:], **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter_sort_static(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sort_static", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_class_sort3", function_name="sort_static", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_Sorter.sort_static(*args, **kwargs) def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="__init__", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_class_sort4", function_name="__init__", file_path=r"{bubble_sort_path}", class_name="Sorter", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) function_name = "__init__" @@ -131,20 +139,29 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): """ assert test_class_sort_path.read_text("utf-8").strip()==test_class_sort_code.strip() - test_sort_path = tests_root / Path("test_benchmark_bubble_sort_test_sort__replay_test_0.py") + test_sort_path = tests_root / Path("test_process_and_sort__replay_test_0.py") assert test_sort_path.exists() test_sort_code = f""" import dill as pickle from code_to_optimize.bubble_sort_codeflash_trace import \\ sorter as code_to_optimize_bubble_sort_codeflash_trace_sorter +from code_to_optimize.process_and_bubble_sort_codeflash_trace import \\ + compute_and_sort as \\ + code_to_optimize_process_and_bubble_sort_codeflash_trace_compute_and_sort from codeflash.benchmarking.replay_test import get_next_arg_and_return -functions = ['sorter'] +functions = ['compute_and_sort', 'sorter'] trace_file_path = r"{output_file}" +def test_code_to_optimize_process_and_bubble_sort_codeflash_trace_compute_and_sort(): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_compute_and_sort", function_name="compute_and_sort", file_path=r"{process_and_bubble_sort_path}", num_to_get=100): + args = pickle.loads(args_pkl) + kwargs = pickle.loads(kwargs_pkl) + ret = code_to_optimize_process_and_bubble_sort_codeflash_trace_compute_and_sort(*args, **kwargs) + def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): - for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, function_name="sorter", file_path=r"{bubble_sort_path}", num_to_get=100): + for args_pkl, kwargs_pkl in get_next_arg_and_return(trace_file=trace_file_path, benchmark_function_name="test_no_func", function_name="sorter", file_path=r"{bubble_sort_path}", num_to_get=100): args = pickle.loads(args_pkl) kwargs = pickle.loads(kwargs_pkl) ret = code_to_optimize_bubble_sort_codeflash_trace_sorter(*args, **kwargs) From b3c83204f8ffea9635f3cafb2d3297cef245c615 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 13:45:38 -0700 Subject: [PATCH 089/122] Use module path instead of file path for benchmarks, improved display to console. --- ... => test_benchmark_bubble_sort_example.py} | 0 ...rt.py => test_process_and_sort_example.py} | 0 codeflash/benchmarking/plugin/plugin.py | 13 ++++-- .../pytest_new_process_trace_benchmarks.py | 4 +- codeflash/benchmarking/replay_test.py | 8 +--- codeflash/benchmarking/trace_benchmarks.py | 9 ++-- codeflash/benchmarking/utils.py | 6 ++- codeflash/models/models.py | 1 - codeflash/optimization/function_optimizer.py | 38 ++++++--------- codeflash/optimization/optimizer.py | 26 ++++------- codeflash/result/explanation.py | 46 +++++++++++++++++-- tests/test_trace_benchmarks.py | 38 +++++++-------- 12 files changed, 106 insertions(+), 83 deletions(-) rename code_to_optimize/tests/pytest/benchmarks_test/{test_benchmark_bubble_sort.py => test_benchmark_bubble_sort_example.py} (100%) rename code_to_optimize/tests/pytest/benchmarks_test/{test_process_and_sort.py => test_process_and_sort_example.py} (100%) diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py b/code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort_example.py similarity index 100% rename from code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort.py rename to code_to_optimize/tests/pytest/benchmarks_test/test_benchmark_bubble_sort_example.py diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py b/code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort_example.py similarity index 100% rename from code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort.py rename to code_to_optimize/tests/pytest/benchmarks_test/test_process_and_sort_example.py diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index b022f9afb..fc19b19d5 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -1,11 +1,15 @@ from __future__ import annotations + import os import sqlite3 import sys import time from pathlib import Path + import pytest + from codeflash.benchmarking.codeflash_trace import codeflash_trace +from codeflash.code_utils.code_utils import module_name_from_file_path from codeflash.models.models import BenchmarkKey @@ -13,11 +17,13 @@ class CodeFlashBenchmarkPlugin: def __init__(self) -> None: self._trace_path = None self._connection = None + self.project_root = None self.benchmark_timings = [] - def setup(self, trace_path:str) -> None: + def setup(self, trace_path:str, project_root:str) -> None: try: # Open connection + self.project_root = project_root self._trace_path = trace_path self._connection = sqlite3.connect(self._trace_path) cur = self._connection.cursor() @@ -235,9 +241,10 @@ def test_something(benchmark): Returns: The return value of the function + a """ - benchmark_file_path = str(self.request.node.fspath) + benchmark_file_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 @@ -273,4 +280,4 @@ def benchmark(request): return CodeFlashBenchmarkPlugin.Benchmark(request) -codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin() \ No newline at end of file +codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin() diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 7b6bd747a..1bb7bbfa4 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py @@ -13,7 +13,7 @@ import pytest try: - codeflash_benchmark_plugin.setup(trace_file) + 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] @@ -22,4 +22,4 @@ except Exception as e: print(f"Failed to collect tests: {e!s}", file=sys.stderr) exitcode = -1 - sys.exit(exitcode) \ No newline at end of file + sys.exit(exitcode) diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 6466b24db..5b654de92 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -271,20 +271,16 @@ def generate_replay_test(trace_file_path: Path, output_dir: Path, test_framework max_run_count=max_run_count, ) test_code = isort.code(test_code) - - name = Path(benchmark_file_path).name.split(".")[0][5:] # remove "test_" from the name since we add it in later output_file = get_test_file_path( - test_dir=Path(output_dir), function_name=f"{name}", test_type="replay" + test_dir=Path(output_dir), function_name=benchmark_file_path, test_type="replay" ) # Write test code to file, parents = true output_dir.mkdir(parents=True, exist_ok=True) output_file.write_text(test_code, "utf-8") count += 1 - logger.info(f"Replay test for benchmark file `{benchmark_file_path}` in {name} written to {output_file}") conn.close() - except Exception as e: logger.info(f"Error generating replay tests: {e}") - return count \ No newline at end of file + return count diff --git a/codeflash/benchmarking/trace_benchmarks.py b/codeflash/benchmarking/trace_benchmarks.py index 8882078d9..8f68030cb 100644 --- a/codeflash/benchmarking/trace_benchmarks.py +++ b/codeflash/benchmarking/trace_benchmarks.py @@ -1,13 +1,12 @@ from __future__ import annotations import re - -from pytest import ExitCode +import subprocess +from pathlib import Path from codeflash.cli_cmds.console import logger from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE -from pathlib import Path -import subprocess + def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root: Path, trace_file: Path, timeout:int = 300) -> None: result = subprocess.run( @@ -40,4 +39,4 @@ def trace_benchmarks_pytest(benchmarks_root: Path, tests_root:Path, project_root error_section = result.stdout logger.warning( f"Error collecting benchmarks - Pytest Exit code: {result.returncode}, {error_section}" - ) \ No newline at end of file + ) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 1d8b22f50..212512ea9 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,4 +1,5 @@ from __future__ import annotations + from typing import Optional from rich.console import Console @@ -6,7 +7,7 @@ from codeflash.cli_cmds.console import logger from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.models.models import ProcessedBenchmarkInfo, BenchmarkDetail, BenchmarkKey +from codeflash.models.models import BenchmarkDetail, BenchmarkKey, ProcessedBenchmarkInfo from codeflash.result.critic import performance_gain @@ -37,6 +38,7 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, di def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey, float, float, float]]]) -> None: console = Console() for func_path, sorted_tests in function_to_results.items(): + console.print() function_name = func_path.split(":")[-1] # Create a table for this function @@ -114,4 +116,4 @@ def process_benchmark_data( ) ) - return ProcessedBenchmarkInfo(benchmark_details=benchmark_details) \ No newline at end of file + return ProcessedBenchmarkInfo(benchmark_details=benchmark_details) diff --git a/codeflash/models/models.py b/codeflash/models/models.py index e415108e7..581c0650b 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -473,7 +473,6 @@ def merge(self, other: TestResults) -> None: def group_by_benchmarks(self, benchmark_keys:list[BenchmarkKey], benchmark_replay_test_dir: Path, project_root: Path) -> dict[BenchmarkKey, TestResults]: """Group TestResults by benchmark for calculating improvements for each benchmark.""" - test_results_by_benchmark = defaultdict(TestResults) benchmark_module_path = {} for benchmark_key in benchmark_keys: diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 807fd3a8c..53a342057 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -42,7 +42,6 @@ from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast from codeflash.code_utils.time_utils import humanize_runtime from codeflash.context import code_context_extractor -from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.either import Failure, Success, is_successful from codeflash.models.ExperimentMetadata import ExperimentMetadata from codeflash.models.models import ( @@ -58,7 +57,7 @@ TestFiles, TestingMode, TestResults, - TestType, BenchmarkKey, + TestType, ) from codeflash.result.create_pr import check_create_pr, existing_tests_source_for from codeflash.result.critic import coverage_critic, performance_gain, quantity_of_tests_critic, speedup_critic @@ -75,8 +74,9 @@ if TYPE_CHECKING: from argparse import Namespace + from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.either import Result - from codeflash.models.models import CoverageData, FunctionSource, OptimizedCandidate + from codeflash.models.models import BenchmarkKey, CoverageData, FunctionSource, OptimizedCandidate from codeflash.verification.verification_utils import TestConfig @@ -92,6 +92,7 @@ def __init__( function_benchmark_timings: dict[BenchmarkKey, int] | None = None, total_benchmark_timings: dict[BenchmarkKey, int] | None = None, args: Namespace | None = None, + replay_tests_dir: Path|None = None ) -> None: self.project_root = test_cfg.project_root_path self.test_cfg = test_cfg @@ -120,6 +121,7 @@ def __init__( self.function_benchmark_timings = function_benchmark_timings if function_benchmark_timings else {} self.total_benchmark_timings = total_benchmark_timings if total_benchmark_timings else {} + self.replay_tests_dir = replay_tests_dir if replay_tests_dir else None def optimize_function(self) -> Result[BestOptimization, str]: should_run_experiment = self.experiment_id is not None @@ -392,7 +394,7 @@ def determine_best_candidate( ) continue - # Instrument codeflash capture + run_results = self.run_optimized_candidate( optimization_candidate_index=candidate_index, baseline_results=original_code_baseline, @@ -430,8 +432,8 @@ def determine_best_candidate( tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") replay_perf_gain = {} if self.args.benchmark: - logger.info(f"Calculating benchmark improvement..") - test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) + benchmark_tree = Tree("Speedup percentage on benchmarks:") + test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.replay_tests_dir, self.project_root) for benchmark_key, candidate_test_results in test_results_by_benchmark.items(): original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results[benchmark_key].total_passed_runtime() candidate_replay_runtime = candidate_test_results.total_passed_runtime() @@ -439,15 +441,8 @@ def determine_best_candidate( original_runtime_ns=original_code_replay_runtime, optimized_runtime_ns=candidate_replay_runtime, ) - tree.add( - f"Original benchmark replay runtime: {humanize_runtime(original_code_replay_runtime)}") - tree.add( - f"Best benchmark replay runtime: {humanize_runtime(candidate_replay_runtime)} " - f"(measured over {candidate_result.max_loop_count} " - f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" - ) - tree.add(f"Speedup percentage for benchmark replay test: {replay_perf_gain[benchmark_key] * 100:.1f}%") - tree.add(f"Speedup ratio for benchmark replay test: {replay_perf_gain[benchmark_key] + 1:.1f}X") + benchmark_tree.add(f"{benchmark_key}: {replay_perf_gain[benchmark_key] * 100:.1f}%") + best_optimization = BestOptimization( candidate=candidate, helper_functions=code_context.helper_functions, @@ -467,6 +462,8 @@ def determine_best_candidate( tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X") console.print(tree) + if self.args.benchmark and benchmark_tree: + console.print(benchmark_tree) console.rule() self.write_code_and_helpers( @@ -903,10 +900,7 @@ def establish_original_code_baseline( logger.debug(f"Total original code runtime (ns): {total_timing}") if self.args.benchmark: - replay_benchmarking_test_results = benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) - for benchmark_name, benchmark_results in replay_benchmarking_test_results.items(): - - logger.info(f"Replay benchmark '{benchmark_name}' runtime: {humanize_runtime(benchmark_results.total_passed_runtime())}") + replay_benchmarking_test_results = benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.replay_tests_dir, self.project_root) return Success( ( OriginalCodeBaseline( @@ -929,7 +923,6 @@ def run_optimized_candidate( file_path_to_helper_classes: dict[Path, set[str]], ) -> Result[OptimizedCandidateResult, str]: assert (test_framework := self.args.test_framework) in ["pytest", "unittest"] - with progress_bar("Testing optimization candidate"): test_env = os.environ.copy() test_env["CODEFLASH_LOOP_INDEX"] = "0" @@ -941,8 +934,6 @@ def run_optimized_candidate( test_env["PYTHONPATH"] += os.pathsep + str(self.project_root) get_run_tmp_file(Path(f"test_return_values_{optimization_candidate_index}.sqlite")).unlink(missing_ok=True) - get_run_tmp_file(Path(f"test_return_values_{optimization_candidate_index}.sqlite")).unlink(missing_ok=True) - # Instrument codeflash capture candidate_fto_code = Path(self.function_to_optimize.file_path).read_text("utf-8") candidate_helper_code = {} @@ -973,7 +964,6 @@ def run_optimized_candidate( ) ) console.rule() - if compare_test_results(baseline_results.behavioral_test_results, candidate_behavior_results): logger.info("Test results matched!") console.rule() @@ -1027,7 +1017,7 @@ def run_optimized_candidate( logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") if self.args.benchmark: - candidate_replay_benchmarking_results = candidate_benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.test_cfg.benchmark_tests_root / "codeflash_replay_tests", self.project_root) + candidate_replay_benchmarking_results = candidate_benchmarking_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.replay_tests_dir, self.project_root) for benchmark_name, benchmark_results in candidate_replay_benchmarking_results.items(): logger.debug(f"Benchmark {benchmark_name} runtime (ns): {humanize_runtime(benchmark_results.total_passed_runtime())}") return Success( diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 51b2ea29e..5e162f513 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -4,10 +4,12 @@ import os import shutil import tempfile +from collections import defaultdict from pathlib import Path from typing import TYPE_CHECKING from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient +from codeflash.benchmarking.instrument_codeflash_trace import instrument_codeflash_trace_decorator from codeflash.benchmarking.plugin.plugin import CodeFlashBenchmarkPlugin from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest @@ -20,16 +22,10 @@ from codeflash.discovery.discover_unit_tests import discover_unit_tests from codeflash.discovery.functions_to_optimize import get_functions_to_optimize from codeflash.either import is_successful -from codeflash.models.models import ValidCode, BenchmarkKey +from codeflash.models.models import BenchmarkKey, TestType, ValidCode from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.telemetry.posthog_cf import ph -from codeflash.verification.test_results import TestType from codeflash.verification.verification_utils import TestConfig -from codeflash.benchmarking.utils import print_benchmark_table -from codeflash.benchmarking.instrument_codeflash_trace import instrument_codeflash_trace_decorator - - -from collections import defaultdict if TYPE_CHECKING: from argparse import Namespace @@ -54,7 +50,7 @@ def __init__(self, args: Namespace) -> None: self.aiservice_client = AiServiceClient() self.experiment_id = os.getenv("CODEFLASH_EXPERIMENT_ID", None) self.local_aiservice_client = LocalAiServiceClient() if self.experiment_id else None - + self.replay_tests_dir = None def create_function_optimizer( self, function_to_optimize: FunctionToOptimize, @@ -74,6 +70,7 @@ def create_function_optimizer( args=self.args, function_benchmark_timings=function_benchmark_timings if function_benchmark_timings else None, total_benchmark_timings=total_benchmark_timings if total_benchmark_timings else None, + replay_tests_dir = self.replay_tests_dir ) def run(self) -> None: @@ -115,9 +112,9 @@ def run(self) -> None: if trace_file.exists(): trace_file.unlink() - replay_tests_dir = Path(self.args.benchmarks_root) / "codeflash_replay_tests" + self.replay_tests_dir = Path(tempfile.mkdtemp(prefix="codeflash_replay_tests_", dir=self.args.benchmarks_root)) trace_benchmarks_pytest(self.args.benchmarks_root, self.args.tests_root, self.args.project_root, trace_file) # Run all tests that use pytest-benchmark - replay_count = generate_replay_test(trace_file, replay_tests_dir) + replay_count = generate_replay_test(trace_file, self.replay_tests_dir) if replay_count == 0: logger.info(f"No valid benchmarks found in {self.args.benchmarks_root} for functions to optimize, continuing optimization") else: @@ -125,10 +122,9 @@ def run(self) -> None: total_benchmark_timings = CodeFlashBenchmarkPlugin.get_benchmark_timings(trace_file) function_to_results = validate_and_format_benchmark_table(function_benchmark_timings, total_benchmark_timings) print_benchmark_table(function_to_results) - logger.info("Finished tracing existing benchmarks") except Exception as e: logger.info(f"Error while tracing existing benchmarks: {e}") - logger.info(f"Information on existing benchmarks will not be available for this run.") + logger.info("Information on existing benchmarks will not be available for this run.") finally: # Restore original source code for file in file_path_to_source_code: @@ -147,8 +143,6 @@ def run(self) -> None: logger.info("No functions found to optimize. Exiting…") return - console.rule() - logger.info(f"Discovering existing unit tests in {self.test_cfg.tests_root}…") console.rule() function_to_tests: dict[str, list[FunctionCalledInTest]] = discover_unit_tests(self.test_cfg) num_discovered_tests: int = sum([len(value) for value in function_to_tests.values()]) @@ -252,8 +246,8 @@ def run(self) -> None: if function_optimizer.test_cfg.concolic_test_root_dir: shutil.rmtree(function_optimizer.test_cfg.concolic_test_root_dir, ignore_errors=True) if self.args.benchmark: - # if replay_tests_dir.exists(): - # shutil.rmtree(replay_tests_dir, ignore_errors=True) + if self.replay_tests_dir.exists(): + shutil.rmtree(self.replay_tests_dir, ignore_errors=True) trace_file.unlink(missing_ok=True) if hasattr(get_run_tmp_file, "tmpdir"): get_run_tmp_file.tmpdir.cleanup() diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index e56558a94..75288bb60 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -1,8 +1,13 @@ from __future__ import annotations + +import shutil +from io import StringIO from pathlib import Path -from typing import Optional, Union +from typing import Optional from pydantic.dataclasses import dataclass +from rich.console import Console +from rich.table import Table from codeflash.code_utils.time_utils import humanize_runtime from codeflash.models.models import BenchmarkDetail, TestResults @@ -43,11 +48,42 @@ def to_console_string(self) -> str: benchmark_info = "" if self.benchmark_details: - benchmark_info += "Benchmark Performance Details:\n" + # Get terminal width (or use a reasonable default if detection fails) + try: + terminal_width = int(shutil.get_terminal_size().columns * 0.8) + except Exception: + terminal_width = 200 # Fallback width + + # Create a rich table for better formatting + table = Table(title="Benchmark Performance Details", width=terminal_width) + + # Add columns - split Benchmark File and Function into separate columns + # Using proportional width for benchmark file column (40% of terminal width) + benchmark_col_width = max(int(terminal_width * 0.4), 40) + table.add_column("Benchmark File", style="cyan", width=benchmark_col_width) + table.add_column("Function", style="cyan") + table.add_column("Original Runtime", style="magenta") + table.add_column("Expected New Runtime", style="green") + table.add_column("Speedup", style="red") + + # Add rows with split data for detail in self.benchmark_details: - benchmark_info += f"Original timing for {detail.benchmark_name}::{detail.test_function}: {detail.original_timing}\n" - benchmark_info += f"Expected new timing for {detail.benchmark_name}::{detail.test_function}: {detail.expected_new_timing}\n" - benchmark_info += f"Benchmark speedup for {detail.benchmark_name}::{detail.test_function}: {detail.speedup_percent:.2f}%\n\n" + # Split the benchmark name and test function + benchmark_name = detail.benchmark_name + test_function = detail.test_function + + table.add_row( + benchmark_name, + test_function, + f"{detail.original_timing}", + f"{detail.expected_new_timing}", + f"{detail.speedup_percent:.2f}%" + ) + + # Render table to string - using actual terminal width + console = Console(file=StringIO(), width=terminal_width) + console.print(table) + benchmark_info = console.file.getvalue() + "\n" return ( f"Optimized {self.function_name} in {self.file_path}\n" diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 7851230d9..c50ab4332 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -13,9 +13,9 @@ 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_test" - tests_root = project_root / "tests" / "test_trace_benchmarks" - tests_root.mkdir(parents=False, exist_ok=False) - output_file = (tests_root / Path("test_trace_benchmarks.trace")).resolve() + replay_tests_dir = benchmarks_root / "codeflash_replay_tests" + 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: @@ -39,31 +39,31 @@ def test_trace_benchmarks(): expected_calls = [ ("sorter", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 17), + "test_class_sort", "tests.pytest.benchmarks_test.test_benchmark_bubble_sort_example", 17), ("sort_class", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort2", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 20), + "test_class_sort2", "tests.pytest.benchmarks_test.test_benchmark_bubble_sort_example", 20), ("sort_static", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort3", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 23), + "test_class_sort3", "tests.pytest.benchmarks_test.test_benchmark_bubble_sort_example", 23), ("__init__", "Sorter", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_class_sort4", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 26), + "test_class_sort4", "tests.pytest.benchmarks_test.test_benchmark_bubble_sort_example", 26), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_sort", str(benchmarks_root / "test_benchmark_bubble_sort.py"), 7), + "test_sort", "tests.pytest.benchmarks_test.test_benchmark_bubble_sort_example", 7), ("compute_and_sort", "", "code_to_optimize.process_and_bubble_sort_codeflash_trace", f"{process_and_bubble_sort_path}", - "test_compute_and_sort", str(benchmarks_root / "test_process_and_sort.py"), 4), + "test_compute_and_sort", "tests.pytest.benchmarks_test.test_process_and_sort_example", 4), ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_no_func", str(benchmarks_root / "test_process_and_sort.py"), 8), + "test_no_func", "tests.pytest.benchmarks_test.test_process_and_sort_example", 8), ] for idx, (actual, expected) in enumerate(zip(function_calls, expected_calls)): assert actual[0] == expected[0], f"Mismatch at index {idx} for function_name" @@ -75,8 +75,8 @@ def test_trace_benchmarks(): assert actual[6] == expected[6], f"Mismatch at index {idx} for benchmark_line_number" # Close connection conn.close() - generate_replay_test(output_file, tests_root) - test_class_sort_path = tests_root / Path("test_benchmark_bubble_sort__replay_test_0.py") + generate_replay_test(output_file, replay_tests_dir) + test_class_sort_path = replay_tests_dir/ Path("test_tests_pytest_benchmarks_test_test_benchmark_bubble_sort_example__replay_test_0.py") assert test_class_sort_path.exists() test_class_sort_code = f""" import dill as pickle @@ -139,7 +139,7 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): """ assert test_class_sort_path.read_text("utf-8").strip()==test_class_sort_code.strip() - test_sort_path = tests_root / Path("test_process_and_sort__replay_test_0.py") + test_sort_path = replay_tests_dir / Path("test_tests_pytest_benchmarks_test_test_process_and_sort_example__replay_test_0.py") assert test_sort_path.exists() test_sort_code = f""" import dill as pickle @@ -170,14 +170,14 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): assert test_sort_path.read_text("utf-8").strip()==test_sort_code.strip() finally: # cleanup - shutil.rmtree(tests_root) + output_file.unlink(missing_ok=True) + shutil.rmtree(replay_tests_dir) def test_trace_multithreaded_benchmark() -> None: project_root = Path(__file__).parent.parent / "code_to_optimize" benchmarks_root = project_root / "tests" / "pytest" / "benchmarks_multithread" - tests_root = project_root / "tests" / "test_trace_benchmarks" - tests_root.mkdir(parents=False, exist_ok=False) - output_file = (tests_root / Path("test_trace_benchmarks.trace")).resolve() + 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: @@ -209,7 +209,7 @@ def test_trace_multithreaded_benchmark() -> None: expected_calls = [ ("sorter", "", "code_to_optimize.bubble_sort_codeflash_trace", f"{bubble_sort_path}", - "test_benchmark_sort", str(benchmarks_root / "test_multithread_sort.py"), 4), + "test_benchmark_sort", "tests.pytest.benchmarks_multithread.test_multithread_sort", 4), ] for idx, (actual, expected) in enumerate(zip(function_calls, expected_calls)): assert actual[0] == expected[0], f"Mismatch at index {idx} for function_name" @@ -224,4 +224,4 @@ def test_trace_multithreaded_benchmark() -> None: finally: # cleanup - shutil.rmtree(tests_root) \ No newline at end of file + output_file.unlink(missing_ok=True) \ No newline at end of file From 972ef4648e72f024e6faa31a28b4564045b11985 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 15:12:19 -0700 Subject: [PATCH 090/122] benchmark flow is working. changed paths to use module_path instead of file_path for Benchmarkkey --- codeflash/benchmarking/codeflash_trace.py | 17 +- codeflash/benchmarking/plugin/plugin.py | 24 +- codeflash/benchmarking/replay_test.py | 14 +- codeflash/benchmarking/utils.py | 26 +- codeflash/models/models.py | 6 +- codeflash/optimization/function_optimizer.py | 295 ++++++++++++------- codeflash/result/explanation.py | 4 +- tests/test_trace_benchmarks.py | 8 +- 8 files changed, 242 insertions(+), 152 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index bcbb3268c..a2d080283 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/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 diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index fc19b19d5..c7c11c6d4 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/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 diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 5b654de92..63a330774 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/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) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 212512ea9..dff32b57e 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,5 +1,6 @@ from __future__ import annotations +import shutil from typing import Optional from rich.console import Console @@ -35,8 +36,14 @@ 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] @@ -44,23 +51,30 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey # 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)), diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 581c0650b..915a05c9a 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -80,11 +80,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: @@ -476,7 +476,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(): diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index 53a342057..da9378f12 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -7,7 +7,7 @@ import subprocess import time import uuid -from collections import defaultdict +from collections import defaultdict, deque from pathlib import Path from typing import TYPE_CHECKING @@ -38,6 +38,7 @@ ) from codeflash.code_utils.formatter import format_code, sort_imports from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test +from codeflash.code_utils.line_profile_utils import add_decorator_imports from codeflash.code_utils.remove_generated_tests import remove_functions_from_generated_tests from codeflash.code_utils.static_analysis import get_first_top_level_function_or_method_ast from codeflash.code_utils.time_utils import humanize_runtime @@ -66,8 +67,9 @@ from codeflash.verification.concolic_testing import generate_concolic_tests from codeflash.verification.equivalence import compare_test_results from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture +from codeflash.verification.parse_line_profile_test_output import parse_line_profile_results from codeflash.verification.parse_test_output import parse_test_results -from codeflash.verification.test_runner import run_behavioral_tests, run_benchmarking_tests +from codeflash.verification.test_runner import run_behavioral_tests, run_benchmarking_tests, run_line_profile_tests from codeflash.verification.verification_utils import get_test_file_path from codeflash.verification.verifier import generate_tests @@ -237,7 +239,8 @@ def optimize_function(self) -> Result[BestOptimization, str]: ): cleanup_paths(paths_to_cleanup) return Failure("The threshold for test coverage was not met.") - + # request for new optimizations but don't block execution, check for completion later + # adding to control and experiment set but with same traceid best_optimization = None for _u, candidates in enumerate([optimizations_set.control, optimizations_set.experiment]): @@ -371,110 +374,141 @@ def determine_best_candidate( f"{self.function_to_optimize.qualified_name}…" ) console.rule() - try: - for candidate_index, candidate in enumerate(candidates, start=1): - get_run_tmp_file(Path(f"test_return_values_{candidate_index}.bin")).unlink(missing_ok=True) - get_run_tmp_file(Path(f"test_return_values_{candidate_index}.sqlite")).unlink(missing_ok=True) - logger.info(f"Optimization candidate {candidate_index}/{len(candidates)}:") - code_print(candidate.source_code) - try: - did_update = self.replace_function_and_helpers_with_optimized_code( - code_context=code_context, optimized_code=candidate.source_code - ) - if not did_update: - logger.warning( - "No functions were replaced in the optimized code. Skipping optimization candidate." + candidates = deque(candidates) + # Start a new thread for AI service request, start loop in main thread + # check if aiservice request is complete, when it is complete, append result to the candidates list + with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor: + future_line_profile_results = executor.submit( + self.aiservice_client.optimize_python_code_line_profiler, + source_code=code_context.read_writable_code, + dependency_code=code_context.read_only_context_code, + trace_id=self.function_trace_id, + line_profiler_results=original_code_baseline.line_profile_results["str_out"], + num_candidates=10, + experiment_metadata=None, + ) + try: + candidate_index = 0 + done = False + original_len = len(candidates) + while candidates: + # for candidate_index, candidate in enumerate(candidates, start=1): + done = True if future_line_profile_results is None else future_line_profile_results.done() + if done and (future_line_profile_results is not None): + line_profile_results = future_line_profile_results.result() + candidates.extend(line_profile_results) + original_len+= len(candidates) + logger.info(f"Added results from line profiler to candidates, total candidates now: {original_len}") + future_line_profile_results = None + candidate_index += 1 + candidate = candidates.popleft() + get_run_tmp_file(Path(f"test_return_values_{candidate_index}.bin")).unlink(missing_ok=True) + get_run_tmp_file(Path(f"test_return_values_{candidate_index}.sqlite")).unlink(missing_ok=True) + logger.info(f"Optimization candidate {candidate_index}/{original_len}:") + code_print(candidate.source_code) + try: + did_update = self.replace_function_and_helpers_with_optimized_code( + code_context=code_context, optimized_code=candidate.source_code + ) + if not did_update: + logger.warning( + "No functions were replaced in the optimized code. Skipping optimization candidate." + ) + console.rule() + continue + except (ValueError, SyntaxError, cst.ParserSyntaxError, AttributeError) as e: + logger.error(e) + self.write_code_and_helpers( + self.function_to_optimize_source_code, + original_helper_code, + self.function_to_optimize.file_path, ) - console.rule() continue - except (ValueError, SyntaxError, cst.ParserSyntaxError, AttributeError) as e: - logger.error(e) - self.write_code_and_helpers( - self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path - ) - continue - - run_results = self.run_optimized_candidate( - optimization_candidate_index=candidate_index, - baseline_results=original_code_baseline, - original_helper_code=original_helper_code, - file_path_to_helper_classes=file_path_to_helper_classes, - ) - console.rule() - if not is_successful(run_results): - optimized_runtimes[candidate.optimization_id] = None - is_correct[candidate.optimization_id] = False - speedup_ratios[candidate.optimization_id] = None - else: - candidate_result: OptimizedCandidateResult = run_results.unwrap() - best_test_runtime = candidate_result.best_test_runtime - optimized_runtimes[candidate.optimization_id] = best_test_runtime - is_correct[candidate.optimization_id] = True - perf_gain = performance_gain( - original_runtime_ns=original_code_baseline.runtime, optimized_runtime_ns=best_test_runtime + run_results = self.run_optimized_candidate( + optimization_candidate_index=candidate_index, + baseline_results=original_code_baseline, + original_helper_code=original_helper_code, + file_path_to_helper_classes=file_path_to_helper_classes, ) - speedup_ratios[candidate.optimization_id] = perf_gain - - tree = Tree(f"Candidate #{candidate_index} - Runtime Information") - if speedup_critic( - candidate_result, original_code_baseline.runtime, best_runtime_until_now - ) and quantity_of_tests_critic(candidate_result): - tree.add("This candidate is faster than the previous best candidate. πŸš€") - tree.add(f"Original summed runtime: {humanize_runtime(original_code_baseline.runtime)}") - tree.add( - f"Best summed runtime: {humanize_runtime(candidate_result.best_test_runtime)} " - f"(measured over {candidate_result.max_loop_count} " - f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" - ) - tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") - replay_perf_gain = {} - if self.args.benchmark: - benchmark_tree = Tree("Speedup percentage on benchmarks:") - test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.replay_tests_dir, self.project_root) - for benchmark_key, candidate_test_results in test_results_by_benchmark.items(): - original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results[benchmark_key].total_passed_runtime() - candidate_replay_runtime = candidate_test_results.total_passed_runtime() - replay_perf_gain[benchmark_key] = performance_gain( - original_runtime_ns=original_code_replay_runtime, - optimized_runtime_ns=candidate_replay_runtime, - ) - benchmark_tree.add(f"{benchmark_key}: {replay_perf_gain[benchmark_key] * 100:.1f}%") - - best_optimization = BestOptimization( - candidate=candidate, - helper_functions=code_context.helper_functions, - runtime=best_test_runtime, - winning_behavioral_test_results=candidate_result.behavior_test_results, - replay_performance_gain=replay_perf_gain if self.args.benchmark else None, - winning_benchmarking_test_results=candidate_result.benchmarking_test_results, - winning_replay_benchmarking_test_results=candidate_result.benchmarking_test_results, - ) - best_runtime_until_now = best_test_runtime + console.rule() + + if not is_successful(run_results): + optimized_runtimes[candidate.optimization_id] = None + is_correct[candidate.optimization_id] = False + speedup_ratios[candidate.optimization_id] = None else: - tree.add( - f"Summed runtime: {humanize_runtime(best_test_runtime)} " - f"(measured over {candidate_result.max_loop_count} " - f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + candidate_result: OptimizedCandidateResult = run_results.unwrap() + best_test_runtime = candidate_result.best_test_runtime + optimized_runtimes[candidate.optimization_id] = best_test_runtime + is_correct[candidate.optimization_id] = True + perf_gain = performance_gain( + original_runtime_ns=original_code_baseline.runtime, optimized_runtime_ns=best_test_runtime + ) + speedup_ratios[candidate.optimization_id] = perf_gain + + tree = Tree(f"Candidate #{candidate_index} - Runtime Information") + if speedup_critic( + candidate_result, original_code_baseline.runtime, best_runtime_until_now + ) and quantity_of_tests_critic(candidate_result): + tree.add("This candidate is faster than the previous best candidate. πŸš€") + tree.add(f"Original summed runtime: {humanize_runtime(original_code_baseline.runtime)}") + tree.add( + f"Best summed runtime: {humanize_runtime(candidate_result.best_test_runtime)} " + f"(measured over {candidate_result.max_loop_count} " + f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + ) + tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio: {perf_gain + 1:.1f}X") + replay_perf_gain = {} + if self.args.benchmark: + test_results_by_benchmark = candidate_result.benchmarking_test_results.group_by_benchmarks(self.total_benchmark_timings.keys(), self.replay_tests_dir, self.project_root) + if len(test_results_by_benchmark) > 0: + benchmark_tree = Tree("Speedup percentage on benchmarks:") + for benchmark_key, candidate_test_results in test_results_by_benchmark.items(): + + original_code_replay_runtime = original_code_baseline.replay_benchmarking_test_results[benchmark_key].total_passed_runtime() + candidate_replay_runtime = candidate_test_results.total_passed_runtime() + replay_perf_gain[benchmark_key] = performance_gain( + original_runtime_ns=original_code_replay_runtime, + optimized_runtime_ns=candidate_replay_runtime, + ) + benchmark_tree.add(f"{benchmark_key}: {replay_perf_gain[benchmark_key] * 100:.1f}%") + + best_optimization = BestOptimization( + candidate=candidate, + helper_functions=code_context.helper_functions, + runtime=best_test_runtime, + winning_behavioral_test_results=candidate_result.behavior_test_results, + replay_performance_gain=replay_perf_gain if self.args.benchmark else None, + winning_benchmarking_test_results=candidate_result.benchmarking_test_results, + winning_replay_benchmarking_test_results=candidate_result.benchmarking_test_results, + ) + best_runtime_until_now = best_test_runtime + else: + tree.add( + f"Summed runtime: {humanize_runtime(best_test_runtime)} " + f"(measured over {candidate_result.max_loop_count} " + f"loop{'s' if candidate_result.max_loop_count > 1 else ''})" + ) + tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") + tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X") + console.print(tree) + if self.args.benchmark and benchmark_tree: + console.print(benchmark_tree) + console.rule() + + self.write_code_and_helpers( + self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path ) - tree.add(f"Speedup percentage: {perf_gain * 100:.1f}%") - tree.add(f"Speedup ratio: {perf_gain + 1:.3f}X") - console.print(tree) - if self.args.benchmark and benchmark_tree: - console.print(benchmark_tree) - console.rule() + except KeyboardInterrupt as e: self.write_code_and_helpers( self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path ) - except KeyboardInterrupt as e: - self.write_code_and_helpers( - self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path - ) - logger.exception(f"Optimization interrupted: {e}") - raise + logger.exception(f"Optimization interrupted: {e}") + raise self.aiservice_client.log_results( function_trace_id=self.function_trace_id, @@ -792,6 +826,7 @@ def establish_original_code_baseline( original_helper_code: dict[Path, str], file_path_to_helper_classes: dict[Path, set[str]], ) -> Result[tuple[OriginalCodeBaseline, list[str]], str]: + line_profile_results = {"timings": {}, "unit": 0, "str_out": ""} # For the original function - run the tests and get the runtime, plus coverage with progress_bar(f"Establishing original code baseline for {self.function_to_optimize.function_name}"): assert (test_framework := self.args.test_framework) in ["pytest", "unittest"] @@ -835,6 +870,28 @@ def establish_original_code_baseline( if not coverage_critic(coverage_results, self.args.test_framework): return Failure("The threshold for test coverage was not met.") if test_framework == "pytest": + try: + line_profiler_output_file = add_decorator_imports(self.function_to_optimize, code_context) + line_profile_results, _ = self.run_and_parse_tests( + testing_type=TestingMode.LINE_PROFILE, + test_env=test_env, + test_files=self.test_files, + optimization_iteration=0, + testing_time=TOTAL_LOOPING_TIME, + enable_coverage=False, + code_context=code_context, + line_profiler_output_file=line_profiler_output_file, + ) + finally: + # Remove codeflash capture + self.write_code_and_helpers( + self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path + ) + if line_profile_results["str_out"] == "": + logger.warning( + f"Couldn't run line profiler for original function {self.function_to_optimize.function_name}" + ) + console.rule() benchmarking_results, _ = self.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, @@ -909,6 +966,7 @@ def establish_original_code_baseline( replay_benchmarking_test_results = replay_benchmarking_test_results if self.args.benchmark else None, runtime=total_timing, coverage_results=coverage_results, + line_profile_results=line_profile_results, ), functions_to_remove, ) @@ -1045,7 +1103,8 @@ def run_and_parse_tests( pytest_max_loops: int = 100_000, code_context: CodeOptimizationContext | None = None, unittest_loop_index: int | None = None, - ) -> tuple[TestResults, CoverageData | None]: + line_profiler_output_file: Path | None = None, + ) -> tuple[TestResults | dict, CoverageData | None]: coverage_database_file = None coverage_config_file = None try: @@ -1059,6 +1118,19 @@ def run_and_parse_tests( verbose=True, enable_coverage=enable_coverage, ) + elif testing_type == TestingMode.LINE_PROFILE: + result_file_path, run_result = run_line_profile_tests( + test_files, + cwd=self.project_root, + test_env=test_env, + pytest_cmd=self.test_cfg.pytest_cmd, + pytest_timeout=INDIVIDUAL_TESTCASE_TIMEOUT, + pytest_target_runtime_seconds=testing_time, + pytest_min_loops=1, + pytest_max_loops=1, + test_framework=self.test_cfg.test_framework, + line_profiler_output_file=line_profiler_output_file, + ) elif testing_type == TestingMode.PERFORMANCE: result_file_path, run_result = run_benchmarking_tests( test_files, @@ -1086,19 +1158,22 @@ def run_and_parse_tests( f"stdout: {run_result.stdout}\n" f"stderr: {run_result.stderr}\n" ) - results, coverage_results = parse_test_results( - test_xml_path=result_file_path, - test_files=test_files, - test_config=self.test_cfg, - optimization_iteration=optimization_iteration, - run_result=run_result, - unittest_loop_index=unittest_loop_index, - function_name=self.function_to_optimize.function_name, - source_file=self.function_to_optimize.file_path, - code_context=code_context, - coverage_database_file=coverage_database_file, - coverage_config_file=coverage_config_file, - ) + if testing_type in [TestingMode.BEHAVIOR, TestingMode.PERFORMANCE]: + results, coverage_results = parse_test_results( + test_xml_path=result_file_path, + test_files=test_files, + test_config=self.test_cfg, + optimization_iteration=optimization_iteration, + run_result=run_result, + unittest_loop_index=unittest_loop_index, + function_name=self.function_to_optimize.function_name, + source_file=self.function_to_optimize.file_path, + code_context=code_context, + coverage_database_file=coverage_database_file, + coverage_config_file=coverage_config_file, + ) + else: + results, coverage_results = parse_line_profile_results(line_profiler_output_file=line_profiler_output_file) return results, coverage_results def generate_and_instrument_tests( diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 75288bb60..076fb0e55 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -60,8 +60,8 @@ def to_console_string(self) -> str: # Add columns - split Benchmark File and Function into separate columns # Using proportional width for benchmark file column (40% of terminal width) benchmark_col_width = max(int(terminal_width * 0.4), 40) - table.add_column("Benchmark File", style="cyan", width=benchmark_col_width) - table.add_column("Function", style="cyan") + table.add_column("Benchmark Module Path", style="cyan", width=benchmark_col_width) + table.add_column("Test Function", style="cyan") table.add_column("Original Runtime", style="magenta") table.add_column("Expected New Runtime", style="green") table.add_column("Speedup", style="red") diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index c50ab4332..e953d1e81 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -27,7 +27,7 @@ def test_trace_benchmarks(): # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_path, benchmark_function_name, benchmark_file_path, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_path, benchmark_function_name, function_name") + "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 @@ -71,7 +71,7 @@ def test_trace_benchmarks(): 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_file_path" + 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() @@ -189,7 +189,7 @@ def test_trace_multithreaded_benchmark() -> None: # Get the count of records # Get all records cursor.execute( - "SELECT function_name, class_name, module_name, file_path, benchmark_function_name, benchmark_file_path, benchmark_line_number FROM benchmark_function_timings ORDER BY benchmark_file_path, benchmark_function_name, function_name") + "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 @@ -217,7 +217,7 @@ def test_trace_multithreaded_benchmark() -> None: 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_file_path" + 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() From 06b3818796f058e15659f8ed7684c1103c594f33 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 15:17:08 -0700 Subject: [PATCH 091/122] fixed string error --- codeflash/models/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 915a05c9a..f19cd660f 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -476,7 +476,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_{benchmark_key.module_path.replace(".", "_")}__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(): From 37577e7cda1c3f4df09b162d719924eaba659922 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 15:36:43 -0700 Subject: [PATCH 092/122] fixed mypy error --- codeflash/result/explanation.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 076fb0e55..217010af2 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -3,7 +3,7 @@ import shutil from io import StringIO from pathlib import Path -from typing import Optional +from typing import Optional, cast from pydantic.dataclasses import dataclass from rich.console import Console @@ -79,11 +79,11 @@ def to_console_string(self) -> str: f"{detail.expected_new_timing}", f"{detail.speedup_percent:.2f}%" ) - - # Render table to string - using actual terminal width - console = Console(file=StringIO(), width=terminal_width) + # Convert table to string + string_buffer = StringIO() + console = Console(file=string_buffer, width=terminal_width) console.print(table) - benchmark_info = console.file.getvalue() + "\n" + benchmark_info = cast(StringIO, console.file).getvalue() + "\n" # Cast for mypy return ( f"Optimized {self.function_name} in {self.file_path}\n" From 5c30d3ea4619a8189adbad722eec2399e9b4d3be Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:15:27 -0700 Subject: [PATCH 093/122] new end to end test for benchmarking bubble sort --- .../workflows/end-to-end-benchmark-test.yaml | 41 +++++++++++++++++++ codeflash/optimization/function_optimizer.py | 7 ++-- .../scripts/end_to_end_test_benchmark_sort.py | 26 ++++++++++++ tests/scripts/end_to_end_test_utilities.py | 8 ++-- 4 files changed, 76 insertions(+), 6 deletions(-) create mode 100644 .github/workflows/end-to-end-benchmark-test.yaml create mode 100644 tests/scripts/end_to_end_test_benchmark_sort.py diff --git a/.github/workflows/end-to-end-benchmark-test.yaml b/.github/workflows/end-to-end-benchmark-test.yaml new file mode 100644 index 000000000..efdb5764f --- /dev/null +++ b/.github/workflows/end-to-end-benchmark-test.yaml @@ -0,0 +1,41 @@ +name: end-to-end-test + +on: + pull_request: + workflow_dispatch: + +jobs: + benchmark-bubble-sort-optimization: + runs-on: ubuntu-latest + env: + CODEFLASH_AIS_SERVER: prod + POSTHOG_API_KEY: ${{ secrets.POSTHOG_API_KEY }} + CODEFLASH_API_KEY: ${{ secrets.CODEFLASH_API_KEY }} + COLUMNS: 110 + MAX_RETRIES: 3 + RETRY_DELAY: 5 + EXPECTED_IMPROVEMENT_PCT: 5 + CODEFLASH_END_TO_END: 1 + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + token: ${{ secrets.GITHUB_TOKEN }} + + - name: Set up Python 3.11 for CLI + uses: astral-sh/setup-uv@v5 + with: + python-version: 3.11.6 + + - name: Install dependencies (CLI) + run: | + uv tool install poetry + uv venv + source .venv/bin/activate + poetry install --with dev + + - name: Run Codeflash to optimize code + id: optimize_code with benchmarks + run: | + source .venv/bin/activate + poetry run python tests/scripts/end_to_end_test_benchmark_sort.py \ No newline at end of file diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py index da9378f12..40b5c7856 100644 --- a/codeflash/optimization/function_optimizer.py +++ b/codeflash/optimization/function_optimizer.py @@ -449,6 +449,7 @@ def determine_best_candidate( speedup_ratios[candidate.optimization_id] = perf_gain tree = Tree(f"Candidate #{candidate_index} - Runtime Information") + benchmark_tree = None if speedup_critic( candidate_result, original_code_baseline.runtime, best_runtime_until_now ) and quantity_of_tests_critic(candidate_result): @@ -499,9 +500,9 @@ def determine_best_candidate( console.print(benchmark_tree) console.rule() - self.write_code_and_helpers( - self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path - ) + self.write_code_and_helpers( + self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path + ) except KeyboardInterrupt as e: self.write_code_and_helpers( diff --git a/tests/scripts/end_to_end_test_benchmark_sort.py b/tests/scripts/end_to_end_test_benchmark_sort.py new file mode 100644 index 000000000..64aabe384 --- /dev/null +++ b/tests/scripts/end_to_end_test_benchmark_sort.py @@ -0,0 +1,26 @@ +import os +import pathlib + +from end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command, run_with_retries + + +def run_test(expected_improvement_pct: int) -> bool: + cwd = (pathlib.Path(__file__).parent.parent.parent / "code_to_optimize").resolve() + config = TestConfig( + file_path=pathlib.Path("bubble_sort.py"), + function_name="sorter", + benchmarks_root=cwd / "tests" / "pytest" / "benchmarks", + test_framework="pytest", + min_improvement_x=1.0, + coverage_expectations=[ + CoverageExpectation( + function_name="sorter", expected_coverage=100.0, expected_lines=[2, 3, 4, 5, 6, 7, 8, 9, 10] + ) + ], + ) + + return run_codeflash_command(cwd, config, expected_improvement_pct) + + +if __name__ == "__main__": + exit(run_with_retries(run_test, int(os.getenv("EXPECTED_IMPROVEMENT_PCT", 5)))) diff --git a/tests/scripts/end_to_end_test_utilities.py b/tests/scripts/end_to_end_test_utilities.py index fda917020..d050f50e9 100644 --- a/tests/scripts/end_to_end_test_utilities.py +++ b/tests/scripts/end_to_end_test_utilities.py @@ -26,6 +26,7 @@ class TestConfig: min_improvement_x: float = 0.1 trace_mode: bool = False coverage_expectations: list[CoverageExpectation] = field(default_factory=list) + benchmarks_root: Optional[pathlib.Path] = None def clear_directory(directory_path: str | pathlib.Path) -> None: @@ -85,8 +86,8 @@ def run_codeflash_command( path_to_file = cwd / config.file_path file_contents = path_to_file.read_text("utf-8") test_root = cwd / "tests" / (config.test_framework or "") - command = build_command(cwd, config, test_root) + command = build_command(cwd, config, test_root, config.benchmarks_root if config.benchmarks_root else None) process = subprocess.Popen( command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, cwd=str(cwd), env=os.environ.copy() ) @@ -116,7 +117,7 @@ def run_codeflash_command( return validated -def build_command(cwd: pathlib.Path, config: TestConfig, test_root: pathlib.Path) -> list[str]: +def build_command(cwd: pathlib.Path, config: TestConfig, test_root: pathlib.Path, benchmarks_root:pathlib.Path|None = None) -> list[str]: python_path = "../../../codeflash/main.py" if "code_directories" in str(cwd) else "../codeflash/main.py" base_command = ["python", python_path, "--file", config.file_path, "--no-pr"] @@ -127,7 +128,8 @@ def build_command(cwd: pathlib.Path, config: TestConfig, test_root: pathlib.Path base_command.extend( ["--test-framework", config.test_framework, "--tests-root", str(test_root), "--module-root", str(cwd)] ) - + if benchmarks_root: + base_command.extend(["--benchmark", "--benchmarks-root", str(benchmarks_root)]) return base_command From 906e4348d142b4e5b9a95a0dfd2f89ef913bcd8c Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:18:18 -0700 Subject: [PATCH 094/122] renamed test --- ...chmark-test.yaml => end-to-end-test-benchmark-bubblesort.yaml} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename .github/workflows/{end-to-end-benchmark-test.yaml => end-to-end-test-benchmark-bubblesort.yaml} (100%) diff --git a/.github/workflows/end-to-end-benchmark-test.yaml b/.github/workflows/end-to-end-test-benchmark-bubblesort.yaml similarity index 100% rename from .github/workflows/end-to-end-benchmark-test.yaml rename to .github/workflows/end-to-end-test-benchmark-bubblesort.yaml From 821fa4798af45bfa0f2be06c30060af5f159795e Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:22:41 -0700 Subject: [PATCH 095/122] fixed e2e test --- .github/workflows/end-to-end-test-benchmark-bubblesort.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/end-to-end-test-benchmark-bubblesort.yaml b/.github/workflows/end-to-end-test-benchmark-bubblesort.yaml index efdb5764f..53a59dac1 100644 --- a/.github/workflows/end-to-end-test-benchmark-bubblesort.yaml +++ b/.github/workflows/end-to-end-test-benchmark-bubblesort.yaml @@ -35,7 +35,7 @@ jobs: poetry install --with dev - name: Run Codeflash to optimize code - id: optimize_code with benchmarks + id: optimize_code_with_benchmarks run: | source .venv/bin/activate poetry run python tests/scripts/end_to_end_test_benchmark_sort.py \ No newline at end of file From 41f7e0a78a3a045ed8e5d765577b4ce7bb9313b7 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:36:47 -0700 Subject: [PATCH 096/122] printing issues on github actions --- codeflash/benchmarking/utils.py | 8 ++++---- codeflash/result/explanation.py | 12 ++++++------ 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index dff32b57e..56644167f 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -40,9 +40,9 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, di def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey, float, float, float]]]) -> None: try: - terminal_width = int(shutil.get_terminal_size().columns * 0.8) + terminal_width = int(shutil.get_terminal_size().columns * 0.9) except Exception: - terminal_width = 200 # Fallback width + terminal_width = 120 # Fallback width console = Console(width = terminal_width) for func_path, sorted_tests in function_to_results.items(): console.print() @@ -52,8 +52,8 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey table = Table(title=f"Function: {function_name}", border_style="blue") # 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("Benchmark Module Path", style="cyan", overflow="fold") + table.add_column("Test Function", style="magenta", overflow="fold") 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") diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 217010af2..2d4aba9bf 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -50,7 +50,7 @@ def to_console_string(self) -> str: if self.benchmark_details: # Get terminal width (or use a reasonable default if detection fails) try: - terminal_width = int(shutil.get_terminal_size().columns * 0.8) + terminal_width = int(shutil.get_terminal_size().columns * 0.9) except Exception: terminal_width = 200 # Fallback width @@ -60,11 +60,11 @@ def to_console_string(self) -> str: # Add columns - split Benchmark File and Function into separate columns # Using proportional width for benchmark file column (40% of terminal width) benchmark_col_width = max(int(terminal_width * 0.4), 40) - table.add_column("Benchmark Module Path", style="cyan", width=benchmark_col_width) - table.add_column("Test Function", style="cyan") - table.add_column("Original Runtime", style="magenta") - table.add_column("Expected New Runtime", style="green") - table.add_column("Speedup", style="red") + table.add_column("Benchmark Module Path", style="cyan", width=benchmark_col_width, overflow="fold") + table.add_column("Test Function", style="cyan", overflow="fold") + table.add_column("Original Runtime", style="magenta", justify="right") + table.add_column("Expected New Runtime", style="green", justify="right") + table.add_column("Speedup", style="red", justify="right") # Add rows with split data for detail in self.benchmark_details: From c20f29aa44bf2d14d2f265b561c1c217daff0736 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:42:29 -0700 Subject: [PATCH 097/122] attempt to use horizontals for rows --- codeflash/benchmarking/utils.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 56644167f..5548feee8 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -1,16 +1,20 @@ from __future__ import annotations import shutil -from typing import Optional +from typing import TYPE_CHECKING, Optional +from rich.box import HORIZONTALS from rich.console import Console from rich.table import Table from codeflash.cli_cmds.console import logger from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.models.models import BenchmarkDetail, BenchmarkKey, ProcessedBenchmarkInfo +from codeflash.models.models import BenchmarkDetail, ProcessedBenchmarkInfo from codeflash.result.critic import performance_gain +if TYPE_CHECKING: + from codeflash.models.models import BenchmarkKey + def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, dict[BenchmarkKey, int]], total_benchmark_timings: dict[BenchmarkKey, int]) -> dict[str, list[tuple[BenchmarkKey, float, float, float]]]: @@ -49,10 +53,10 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey function_name = func_path.split(":")[-1] # Create a table for this function - table = Table(title=f"Function: {function_name}", border_style="blue") - + table = Table(title=f"Function: {function_name}", width=terminal_width, border_style="blue", box=HORIZONTALS) + benchmark_col_width = max(int(terminal_width * 0.4), 40) # Add columns - split the benchmark test into two columns - table.add_column("Benchmark Module Path", style="cyan", overflow="fold") + table.add_column("Benchmark Module Path", width=benchmark_col_width, style="cyan", overflow="fold") table.add_column("Test Function", style="magenta", overflow="fold") table.add_column("Total Time (ms)", justify="right", style="green") table.add_column("Function Time (ms)", justify="right", style="yellow") From d1a8d2524334488d56de57a8cd6aad8ce976e724 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 2 Apr 2025 16:48:10 -0700 Subject: [PATCH 098/122] added row lines --- codeflash/benchmarking/utils.py | 3 +-- codeflash/result/explanation.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index 5548feee8..a10b4487d 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -3,7 +3,6 @@ import shutil from typing import TYPE_CHECKING, Optional -from rich.box import HORIZONTALS from rich.console import Console from rich.table import Table @@ -53,7 +52,7 @@ def print_benchmark_table(function_to_results: dict[str, list[tuple[BenchmarkKey function_name = func_path.split(":")[-1] # Create a table for this function - table = Table(title=f"Function: {function_name}", width=terminal_width, border_style="blue", box=HORIZONTALS) + table = Table(title=f"Function: {function_name}", width=terminal_width, border_style="blue", show_lines=True) benchmark_col_width = max(int(terminal_width * 0.4), 40) # Add columns - split the benchmark test into two columns table.add_column("Benchmark Module Path", width=benchmark_col_width, style="cyan", overflow="fold") diff --git a/codeflash/result/explanation.py b/codeflash/result/explanation.py index 2d4aba9bf..c6e1fb9dc 100644 --- a/codeflash/result/explanation.py +++ b/codeflash/result/explanation.py @@ -55,7 +55,7 @@ def to_console_string(self) -> str: terminal_width = 200 # Fallback width # Create a rich table for better formatting - table = Table(title="Benchmark Performance Details", width=terminal_width) + table = Table(title="Benchmark Performance Details", width=terminal_width, show_lines=True) # Add columns - split Benchmark File and Function into separate columns # Using proportional width for benchmark file column (40% of terminal width) From 705105ce822a3b6bf1356bad6b939c884b80538b Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 3 Apr 2025 15:28:44 -0700 Subject: [PATCH 099/122] made benchmarks-root use resolve() --- codeflash/cli_cmds/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/cli_cmds/cli.py b/codeflash/cli_cmds/cli.py index 07652f707..ed0dbd760 100644 --- a/codeflash/cli_cmds/cli.py +++ b/codeflash/cli_cmds/cli.py @@ -135,7 +135,7 @@ def process_pyproject_config(args: Namespace) -> Namespace: if args.benchmark: assert args.benchmarks_root is not None, "--benchmarks-root must be specified when running with --benchmark" assert Path(args.benchmarks_root).is_dir(), f"--benchmarks-root {args.benchmarks_root} must be a valid directory" - assert Path(args.benchmarks_root).is_relative_to(Path(args.tests_root)), ( + assert Path(args.benchmarks_root).resolve().is_relative_to(Path(args.tests_root).resolve()), ( f"--benchmarks-root {args.benchmarks_root} must be a subdirectory of --tests-root {args.tests_root}" ) if env_utils.get_pr_number() is not None: From 26546de981d8a72fb5d66bbef12895b63a31708d Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 3 Apr 2025 15:59:32 -0700 Subject: [PATCH 100/122] handled edge case for instrumenting codeflash trace --- .../instrument_codeflash_trace.py | 15 +++- tests/test_instrument_codeflash_trace.py | 80 ++++++++++++++++++- 2 files changed, 91 insertions(+), 4 deletions(-) diff --git a/codeflash/benchmarking/instrument_codeflash_trace.py b/codeflash/benchmarking/instrument_codeflash_trace.py index 06e93daf8..044b0b0a4 100644 --- a/codeflash/benchmarking/instrument_codeflash_trace.py +++ b/codeflash/benchmarking/instrument_codeflash_trace.py @@ -12,12 +12,14 @@ def __init__(self, target_functions: set[tuple[str, str]]) -> None: self.target_functions = target_functions self.added_codeflash_trace = False self.class_name = "" + self.function_name = "" self.decorator = cst.Decorator( decorator=cst.Name(value="codeflash_trace") ) def leave_ClassDef(self, original_node, updated_node): - self.class_name = "" + if self.class_name == original_node.name.value: + self.class_name = "" # Even if nested classes are not visited, this function is still called on them return updated_node def visit_ClassDef(self, node): @@ -25,7 +27,14 @@ def visit_ClassDef(self, node): return False self.class_name = node.name.value + def visit_FunctionDef(self, node): + if self.function_name: # Don't go into nested function + return False + self.function_name = node.name.value + def leave_FunctionDef(self, original_node, updated_node): + if self.function_name == original_node.name.value: + self.function_name = "" if (self.class_name, original_node.name.value) in self.target_functions: # Add the new decorator after any existing decorators, so it gets executed first updated_decorators = list(updated_node.decorators) + [self.decorator] @@ -33,8 +42,8 @@ def leave_FunctionDef(self, original_node, updated_node): return updated_node.with_changes( decorators=updated_decorators ) - else: - return updated_node + + return updated_node def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module: # Create import statement for codeflash_trace diff --git a/tests/test_instrument_codeflash_trace.py b/tests/test_instrument_codeflash_trace.py index 6b884c631..38a6381e2 100644 --- a/tests/test_instrument_codeflash_trace.py +++ b/tests/test_instrument_codeflash_trace.py @@ -466,4 +466,82 @@ def static_method_b(): # Compare the modified content with expected content assert modified_content_1.strip() == expected_content_1.strip() - assert modified_content_2.strip() == expected_content_2.strip() \ No newline at end of file + assert modified_content_2.strip() == expected_content_2.strip() + + +def test_add_decorator_to_method_after_nested_class() -> None: + """Test adding decorator to a method that appears after a nested class definition.""" + code = """ +class OuterClass: + class NestedClass: + def nested_method(self): + return "Hello from nested class method" + + def target_method(self): + return "Hello from target method after nested class" +""" + + fto = FunctionToOptimize( + function_name="target_method", + file_path=Path("dummy_path.py"), + parents=[FunctionParent(name="OuterClass", type="ClassDef")] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + functions_to_optimize=[fto] + ) + + expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +class OuterClass: + class NestedClass: + def nested_method(self): + return "Hello from nested class method" + + @codeflash_trace + def target_method(self): + return "Hello from target method after nested class" +""" + + assert modified_code.strip() == expected_code.strip() + + +def test_add_decorator_to_function_after_nested_function() -> None: + """Test adding decorator to a function that appears after a function with a nested function.""" + code = """ +def function_with_nested(): + def inner_function(): + return "Hello from inner function" + + return inner_function() + +def target_function(): + return "Hello from target function after nested function" +""" + + fto = FunctionToOptimize( + function_name="target_function", + file_path=Path("dummy_path.py"), + parents=[] + ) + + modified_code = add_codeflash_decorator_to_code( + code=code, + functions_to_optimize=[fto] + ) + + expected_code = """ +from codeflash.benchmarking.codeflash_trace import codeflash_trace +def function_with_nested(): + def inner_function(): + return "Hello from inner function" + + return inner_function() + +@codeflash_trace +def target_function(): + return "Hello from target function after nested function" +""" + + assert modified_code.strip() == expected_code.strip() \ No newline at end of file From 0c04adf0c75af876ebc3cae779ea9b88a49d4468 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 3 Apr 2025 16:33:41 -0700 Subject: [PATCH 101/122] fixed slight bug with formatting table --- codeflash/benchmarking/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/benchmarking/utils.py b/codeflash/benchmarking/utils.py index a10b4487d..da09cd57a 100644 --- a/codeflash/benchmarking/utils.py +++ b/codeflash/benchmarking/utils.py @@ -29,7 +29,7 @@ def validate_and_format_benchmark_table(function_benchmark_timings: dict[str, di # If the function time is greater than total time, likely to have multithreading / multiprocessing issues. # Do not try to project the optimization impact for this function. sorted_tests.append((benchmark_key, 0.0, 0.0, 0.0)) - if total_time > 0: + elif total_time > 0: percentage = (func_time / total_time) * 100 # Convert nanoseconds to milliseconds func_time_ms = func_time / 1_000_000 From 30d32bbf18321fe28978fd06f01d8b4e5f4df40c Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 3 Apr 2025 16:48:56 -0700 Subject: [PATCH 102/122] improved file removal after errors --- codeflash/optimization/optimizer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 5e162f513..a1260dfd8 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -127,10 +127,11 @@ def run(self) -> None: logger.info("Information on existing benchmarks will not be available for this run.") finally: # Restore original source code + trace_file.unlink() + shutil.rmtree(self.replay_tests_dir, ignore_errors=True) for file in file_path_to_source_code: with file.open("w", encoding="utf8") as f: f.write(file_path_to_source_code[file]) - optimizations_found: int = 0 function_iterator_count: int = 0 if self.args.test_framework == "pytest": From c997b90394406ed4b94cf23b935c124242b6b7fd Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 4 Apr 2025 11:04:58 -0700 Subject: [PATCH 103/122] fixed a return bug --- codeflash/benchmarking/codeflash_trace.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index a2d080283..776e0e635 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -144,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 None + return result if len(self.function_calls_data) > 1000: self.write_function_timings() From d6ed1c33c4a307bbf7ae3be57d22dc6ed25951cb Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Mon, 7 Apr 2025 14:52:18 -0700 Subject: [PATCH 104/122] Support recursive functions, and @benchmark / @pytest.mark.benchmark ways of using benchmark. created tests for all of them --- .../bubble_sort_codeflash_trace.py | 18 +++++ .../benchmarks_test/test_recursive_example.py | 6 ++ .../test_benchmark_decorator.py | 11 +++ codeflash/benchmarking/codeflash_trace.py | 79 +++++++++++++------ codeflash/benchmarking/plugin/plugin.py | 71 ++++++++++------- .../pytest_new_process_trace_benchmarks.py | 2 +- codeflash/benchmarking/replay_test.py | 2 +- tests/test_trace_benchmarks.py | 62 ++++++++++++++- 8 files changed, 194 insertions(+), 57 deletions(-) create mode 100644 code_to_optimize/tests/pytest/benchmarks_test/test_recursive_example.py create mode 100644 code_to_optimize/tests/pytest/benchmarks_test_decorator/test_benchmark_decorator.py diff --git a/code_to_optimize/bubble_sort_codeflash_trace.py b/code_to_optimize/bubble_sort_codeflash_trace.py index ee4dbd999..48e9a412b 100644 --- a/code_to_optimize/bubble_sort_codeflash_trace.py +++ b/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): diff --git a/code_to_optimize/tests/pytest/benchmarks_test/test_recursive_example.py b/code_to_optimize/tests/pytest/benchmarks_test/test_recursive_example.py new file mode 100644 index 000000000..689b1f9ff --- /dev/null +++ b/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)) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/benchmarks_test_decorator/test_benchmark_decorator.py b/code_to_optimize/tests/pytest/benchmarks_test_decorator/test_benchmark_decorator.py new file mode 100644 index 000000000..b924bee7f --- /dev/null +++ b/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)))) \ No newline at end of file diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 776e0e635..95318a38a 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/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, diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index c7c11c6d4..f1614b5c8 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/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,6 +211,13 @@ 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"): @@ -217,9 +225,19 @@ def pytest_collection_modifyitems(config, items): 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: @@ -227,44 +245,37 @@ 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( diff --git a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py b/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py index 1bb7bbfa4..232c39fa7 100644 --- a/codeflash/benchmarking/pytest_new_process_trace_benchmarks.py +++ b/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: diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 63a330774..445957505 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/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: diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index e953d1e81..715955063 100644 --- a/tests/test_trace_benchmarks.py +++ b/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) \ No newline at end of file From 3158f9cc1cf908d063c1cad711b34bd65ea096bb Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 10 Apr 2025 21:43:56 -0400 Subject: [PATCH 105/122] end to end test that proves picklepatcher works. example shown is a socket (which is unpickleable) that's used or not used --- ...ble_sort_picklepatch_test_unused_socket.py | 18 + ...bble_sort_picklepatch_test_used_socket.py} | 36 +- .../benchmarks_socket_test/test_socket.py | 20 ++ .../pytest/test_bubble_sort_picklepatch.py | 34 -- codeflash/benchmarking/codeflash_trace.py | 43 +-- codeflash/benchmarking/plugin/plugin.py | 5 +- codeflash/benchmarking/replay_test.py | 2 +- codeflash/models/models.py | 7 +- codeflash/picklepatch/pickle_placeholder.py | 19 +- codeflash/verification/comparator.py | 7 +- codeflash/verification/parse_test_output.py | 6 +- tests/test_pickle_patcher.py | 324 ++++++++++++++---- tests/test_trace_benchmarks.py | 19 +- 13 files changed, 349 insertions(+), 191 deletions(-) create mode 100644 code_to_optimize/bubble_sort_picklepatch_test_unused_socket.py rename code_to_optimize/{bubble_sort_picklepatch.py => bubble_sort_picklepatch_test_used_socket.py} (55%) create mode 100644 code_to_optimize/tests/pytest/benchmarks_socket_test/test_socket.py delete mode 100644 code_to_optimize/tests/pytest/test_bubble_sort_picklepatch.py diff --git a/code_to_optimize/bubble_sort_picklepatch_test_unused_socket.py b/code_to_optimize/bubble_sort_picklepatch_test_unused_socket.py new file mode 100644 index 000000000..2b75a8c34 --- /dev/null +++ b/code_to_optimize/bubble_sort_picklepatch_test_unused_socket.py @@ -0,0 +1,18 @@ + +from codeflash.benchmarking.codeflash_trace import codeflash_trace + + +@codeflash_trace +def bubble_sort_with_unused_socket(data_container): + # Extract the list to sort, leaving the socket untouched + numbers = data_container.get('numbers', []).copy() + + return sorted(numbers) + +@codeflash_trace +def bubble_sort_with_used_socket(data_container): + # Extract the list to sort, leaving the socket untouched + numbers = data_container.get('numbers', []).copy() + socket = data_container.get('socket') + socket.send("Hello from the optimized function!") + return sorted(numbers) diff --git a/code_to_optimize/bubble_sort_picklepatch.py b/code_to_optimize/bubble_sort_picklepatch_test_used_socket.py similarity index 55% rename from code_to_optimize/bubble_sort_picklepatch.py rename to code_to_optimize/bubble_sort_picklepatch_test_used_socket.py index 25cbe9628..390e090cd 100644 --- a/code_to_optimize/bubble_sort_picklepatch.py +++ b/code_to_optimize/bubble_sort_picklepatch_test_used_socket.py @@ -1,38 +1,6 @@ -def bubble_sort_with_unused_socket(data_container): - """ - Performs a bubble sort on a list within the data_container. The data container has the following schema: - - 'numbers' (list): The list to be sorted. - - 'socket' (socket): A socket - - Args: - data_container: A dictionary with at least 'numbers' (list) and 'socket' keys - - Returns: - list: The sorted list of numbers - """ - # Extract the list to sort, leaving the socket untouched - numbers = data_container.get('numbers', []).copy() - - # Classic bubble sort implementation - n = len(numbers) - for i in range(n): - # Flag to optimize by detecting if no swaps occurred - swapped = False - - # Last i elements are already in place - for j in range(0, n - i - 1): - # Swap if the element is greater than the next element - if numbers[j] > numbers[j + 1]: - numbers[j], numbers[j + 1] = numbers[j + 1], numbers[j] - swapped = True - - # If no swapping occurred in this pass, the list is sorted - if not swapped: - break - - return numbers - +from codeflash.benchmarking.codeflash_trace import codeflash_trace +@codeflash_trace def bubble_sort_with_used_socket(data_container): """ Performs a bubble sort on a list within the data_container. The data container has the following schema: diff --git a/code_to_optimize/tests/pytest/benchmarks_socket_test/test_socket.py b/code_to_optimize/tests/pytest/benchmarks_socket_test/test_socket.py new file mode 100644 index 000000000..bd05af487 --- /dev/null +++ b/code_to_optimize/tests/pytest/benchmarks_socket_test/test_socket.py @@ -0,0 +1,20 @@ +import socket + +from code_to_optimize.bubble_sort_picklepatch_test_unused_socket import bubble_sort_with_unused_socket +from code_to_optimize.bubble_sort_picklepatch_test_used_socket import bubble_sort_with_used_socket + +def test_socket_picklepatch(benchmark): + s1, s2 = socket.socketpair() + data = { + "numbers": list(reversed(range(500))), + "socket": s1 + } + benchmark(bubble_sort_with_unused_socket, data) + +def test_used_socket_picklepatch(benchmark): + s1, s2 = socket.socketpair() + data = { + "numbers": list(reversed(range(500))), + "socket": s1 + } + benchmark(bubble_sort_with_used_socket, data) \ No newline at end of file diff --git a/code_to_optimize/tests/pytest/test_bubble_sort_picklepatch.py b/code_to_optimize/tests/pytest/test_bubble_sort_picklepatch.py deleted file mode 100644 index 9f3e0f9af..000000000 --- a/code_to_optimize/tests/pytest/test_bubble_sort_picklepatch.py +++ /dev/null @@ -1,34 +0,0 @@ -import socket -from unittest.mock import Mock - -import pytest - -from code_to_optimize.bubble_sort_picklepatch import bubble_sort_with_unused_socket, bubble_sort_with_used_socket - - -def test_bubble_sort_with_unused_socket(): - mock_socket = Mock() - # Test case 1: Regular unsorted list - data_container = { - 'numbers': [5, 2, 9, 1, 5, 6], - 'socket': mock_socket - } - - result = bubble_sort_with_unused_socket(data_container) - - # Check that the result is correctly sorted - assert result == [1, 2, 5, 5, 6, 9] - -def test_bubble_sort_with_used_socket(): - mock_socket = Mock() - # Test case 1: Regular unsorted list - data_container = { - 'numbers': [5, 2, 9, 1, 5, 6], - 'socket': mock_socket - } - - result = bubble_sort_with_used_socket(data_container) - - # Check that the result is correctly sorted - assert result == [1, 2, 5, 5, 6, 9] - diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 95318a38a..2694532f3 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -2,12 +2,11 @@ import os import pickle import sqlite3 -import sys import threading import time from typing import Callable -import dill +from codeflash.picklepatch.pickle_patcher import PicklePatcher class CodeflashTrace: @@ -147,34 +146,20 @@ def wrapper(*args, **kwargs): 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: - 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}") - # 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 - + pickled_args = PicklePatcher.dumps(args, protocol=pickle.HIGHEST_PROTOCOL) + pickled_kwargs = PicklePatcher.dumps(kwargs, protocol=pickle.HIGHEST_PROTOCOL) + except Exception 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() diff --git a/codeflash/benchmarking/plugin/plugin.py b/codeflash/benchmarking/plugin/plugin.py index f1614b5c8..313817041 100644 --- a/codeflash/benchmarking/plugin/plugin.py +++ b/codeflash/benchmarking/plugin/plugin.py @@ -175,7 +175,6 @@ 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: @@ -267,9 +266,9 @@ def _run_benchmark(self, func, *args, **kwargs): os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(line_number) os.environ["CODEFLASH_BENCHMARKING"] = "True" # Run the function - start = time.thread_time_ns() + start = time.time_ns() result = func(*args, **kwargs) - end = time.thread_time_ns() + end = time.time_ns() # Reset the environment variable os.environ["CODEFLASH_BENCHMARKING"] = "False" diff --git a/codeflash/benchmarking/replay_test.py b/codeflash/benchmarking/replay_test.py index 445957505..ee1107241 100644 --- a/codeflash/benchmarking/replay_test.py +++ b/codeflash/benchmarking/replay_test.py @@ -62,7 +62,7 @@ def create_trace_replay_test_code( assert test_framework in ["pytest", "unittest"] # Create Imports - imports = f"""import dill as pickle + imports = f"""from codeflash.picklepatch.pickle_patcher import PicklePatcher as pickle {"import unittest" if test_framework == "unittest" else ""} from codeflash.benchmarking.replay_test import get_next_arg_and_return """ diff --git a/codeflash/models/models.py b/codeflash/models/models.py index aede322a1..791912b8a 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -16,7 +16,7 @@ from enum import Enum, IntEnum from pathlib import Path from re import Pattern -from typing import Annotated, Any, Optional, Union, cast +from typing import Annotated, Optional, cast from jedi.api.classes import Name from pydantic import AfterValidator, BaseModel, ConfigDict, Field @@ -362,6 +362,7 @@ class FunctionCoverage: class TestingMode(enum.Enum): BEHAVIOR = "behavior" PERFORMANCE = "performance" + LINE_PROFILE = "line_profile" class VerificationType(str, Enum): @@ -533,7 +534,7 @@ def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree: tree.add( f"{test_type.to_name()} - Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']}" ) - return + return tree def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: @@ -606,4 +607,4 @@ def __eq__(self, other: object) -> bool: sys.setrecursionlimit(original_recursion_limit) return False sys.setrecursionlimit(original_recursion_limit) - return True \ No newline at end of file + return True diff --git a/codeflash/picklepatch/pickle_placeholder.py b/codeflash/picklepatch/pickle_placeholder.py index cddb6535a..a422abb45 100644 --- a/codeflash/picklepatch/pickle_placeholder.py +++ b/codeflash/picklepatch/pickle_placeholder.py @@ -1,3 +1,8 @@ +class PicklePlaceholderAccessError(Exception): + """Custom exception raised when attempting to access an unpicklable object.""" + + + class PicklePlaceholder: """A placeholder for an object that couldn't be pickled. @@ -22,22 +27,22 @@ def __init__(self, obj_type, obj_str, error_msg, path=None): self.__dict__["path"] = path if path is not None else [] def __getattr__(self, name): - """Raise an error when any attribute is accessed.""" + """Raise a custom error when any attribute is accessed.""" path_str = ".".join(self.__dict__["path"]) if self.__dict__["path"] else "root object" - raise AttributeError( - f"Cannot access attribute '{name}' on unpicklable object at {path_str}. " + raise PicklePlaceholderAccessError( + f"Attempt to access unpickleable object: Cannot access attribute '{name}' on unpicklable object at {path_str}. " f"Original type: {self.__dict__['obj_type']}. Error: {self.__dict__['error_msg']}" ) def __setattr__(self, name, value): """Prevent setting attributes.""" - self.__getattr__(name) # This will raise an AttributeError + self.__getattr__(name) # This will raise our custom error def __call__(self, *args, **kwargs): - """Raise an error when the object is called.""" + """Raise a custom error when the object is called.""" path_str = ".".join(self.__dict__["path"]) if self.__dict__["path"] else "root object" - raise TypeError( - f"Cannot call unpicklable object at {path_str}. " + raise PicklePlaceholderAccessError( + f"Attempt to access unpickleable object: Cannot call unpicklable object at {path_str}. " f"Original type: {self.__dict__['obj_type']}. Error: {self.__dict__['error_msg']}" ) diff --git a/codeflash/verification/comparator.py b/codeflash/verification/comparator.py index f047d5b3c..8a7048c57 100644 --- a/codeflash/verification/comparator.py +++ b/codeflash/verification/comparator.py @@ -10,6 +10,7 @@ import sentry_sdk from codeflash.cli_cmds.console import logger +from codeflash.picklepatch.pickle_placeholder import PicklePlaceholderAccessError try: import numpy as np @@ -64,7 +65,11 @@ def comparator(orig: Any, new: Any, superset_obj=False) -> bool: if len(orig) != len(new): return False return all(comparator(elem1, elem2, superset_obj) for elem1, elem2 in zip(orig, new)) - + if isinstance(orig, PicklePlaceholderAccessError) or isinstance(new, PicklePlaceholderAccessError): + # If this error was raised, there was an attempt to access the PicklePlaceholder, which represents an unpickleable object. + # The test results should be rejected as the behavior of the unpickleable object is unknown. + logger.debug("Unable to verify behavior of unpickleable object in replay test") + return False if isinstance( orig, ( diff --git a/codeflash/verification/parse_test_output.py b/codeflash/verification/parse_test_output.py index 80d711894..2228559f9 100644 --- a/codeflash/verification/parse_test_output.py +++ b/codeflash/verification/parse_test_output.py @@ -8,6 +8,7 @@ from pathlib import Path from typing import TYPE_CHECKING +import dill as pickle from junitparser.xunit2 import JUnitXml from lxml.etree import XMLParser, parse @@ -20,7 +21,6 @@ ) from codeflash.discovery.discover_unit_tests import discover_parameters_unittest from codeflash.models.models import FunctionTestInvocation, InvocationId, TestResults, TestType, VerificationType -from codeflash.picklepatch.pickle_patcher import PicklePatcher from codeflash.verification.coverage_utils import CoverageUtils if TYPE_CHECKING: @@ -75,7 +75,7 @@ def parse_test_return_values_bin(file_location: Path, test_files: TestFiles, tes test_type = test_files.get_test_type_by_instrumented_file_path(test_file_path) try: - test_pickle = PicklePatcher.loads(test_pickle_bin) if loop_index == 1 else None + test_pickle = pickle.loads(test_pickle_bin) if loop_index == 1 else None except Exception as e: if DEBUG_MODE: logger.exception(f"Failed to load pickle file for {encoded_test_name} Exception: {e}") @@ -133,7 +133,7 @@ def parse_sqlite_test_results(sqlite_file_path: Path, test_files: TestFiles, tes # TODO : this is because sqlite writes original file module path. Should make it consistent test_type = test_files.get_test_type_by_original_file_path(test_file_path) try: - ret_val = (PicklePatcher.loads(val[7]) if loop_index == 1 else None,) + ret_val = (pickle.loads(val[7]) if loop_index == 1 else None,) except Exception: continue test_results.add( diff --git a/tests/test_pickle_patcher.py b/tests/test_pickle_patcher.py index 05bd06f15..3d2f21b66 100644 --- a/tests/test_pickle_patcher.py +++ b/tests/test_pickle_patcher.py @@ -1,34 +1,40 @@ import os import pickle +import shutil import socket +import sqlite3 from argparse import Namespace from pathlib import Path import dill import pytest -import requests -import sqlite3 +from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin +from codeflash.benchmarking.replay_test import generate_replay_test +from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest +from codeflash.benchmarking.utils import validate_and_format_benchmark_table from codeflash.code_utils.instrument_existing_tests import inject_profiling_into_existing_test from codeflash.discovery.functions_to_optimize import FunctionToOptimize -from codeflash.models.models import CodePosition, TestingMode, TestType, TestFiles, TestFile +from codeflash.models.models import CodePosition, TestFile, TestFiles, TestingMode, TestsInFile, TestType from codeflash.optimization.optimizer import Optimizer +from codeflash.verification.equivalence import compare_test_results try: import sqlalchemy - from sqlalchemy.orm import Session - from sqlalchemy import create_engine, Column, Integer, String + from sqlalchemy import Column, Integer, String, create_engine from sqlalchemy.ext.declarative import declarative_base + from sqlalchemy.orm import Session HAS_SQLALCHEMY = True except ImportError: HAS_SQLALCHEMY = False from codeflash.picklepatch.pickle_patcher import PicklePatcher -from codeflash.picklepatch.pickle_placeholder import PicklePlaceholder +from codeflash.picklepatch.pickle_placeholder import PicklePlaceholder, PicklePlaceholderAccessError + + def test_picklepatch_simple_nested(): - """ - Test that a simple nested data structure pickles and unpickles correctly. + """Test that a simple nested data structure pickles and unpickles correctly. """ original_data = { "numbers": [1, 2, 3], @@ -41,17 +47,24 @@ def test_picklepatch_simple_nested(): assert reloaded == original_data # Everything was pickleable, so no placeholders should appear. + def test_picklepatch_with_socket(): - """ - Test that a data structure containing a raw socket is replaced by + """Test that a data structure containing a raw socket is replaced by PicklePlaceholder rather than raising an error. """ - s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + # Create a pair of connected sockets instead of a single socket + sock1, sock2 = socket.socketpair() + data_with_socket = { "safe_value": 123, - "raw_socket": s, + "raw_socket": sock1, } + # Send a message through sock1, which can be received by sock2 + sock1.send(b"Hello, world!") + received = sock2.recv(1024) + assert received == b"Hello, world!" + # Pickle the data structure containing the socket dumped = PicklePatcher.dumps(data_with_socket) reloaded = PicklePatcher.loads(dumped) @@ -60,15 +73,18 @@ def test_picklepatch_with_socket(): assert reloaded["safe_value"] == 123 assert isinstance(reloaded["raw_socket"], PicklePlaceholder) - # Attempting to use or access attributes => AttributeError + # Attempting to use or access attributes => AttributeError # (not RuntimeError as in original tests, our implementation uses AttributeError) - with pytest.raises(AttributeError) : + with pytest.raises(PicklePlaceholderAccessError): reloaded["raw_socket"].recv(1024) + # Clean up by closing both sockets + sock1.close() + sock2.close() + def test_picklepatch_deeply_nested(): - """ - Test that deep nesting with unpicklable objects works correctly. + """Test that deep nesting with unpicklable objects works correctly. """ # Create a deeply nested structure with an unpicklable object deep_nested = { @@ -92,8 +108,7 @@ def test_picklepatch_deeply_nested(): assert isinstance(reloaded["level1"]["level2"]["level3"]["socket"], PicklePlaceholder) def test_picklepatch_class_with_unpicklable_attr(): - """ - Test that a class with an unpicklable attribute works correctly. + """Test that a class with an unpicklable attribute works correctly. """ class TestClass: def __init__(self): @@ -115,12 +130,11 @@ def __init__(self): def test_picklepatch_with_database_connection(): - """ - Test that a data structure containing a database connection is replaced + """Test that a data structure containing a database connection is replaced by PicklePlaceholder rather than raising an error. """ # SQLite connection - not pickleable - conn = sqlite3.connect(':memory:') + conn = sqlite3.connect(":memory:") cursor = conn.cursor() data_with_db = { @@ -139,13 +153,12 @@ def test_picklepatch_with_database_connection(): assert isinstance(reloaded["cursor"], PicklePlaceholder) # Attempting to use attributes => AttributeError - with pytest.raises(AttributeError): + with pytest.raises(PicklePlaceholderAccessError): reloaded["connection"].execute("SELECT 1") def test_picklepatch_with_generator(): - """ - Test that a data structure containing a generator is replaced by + """Test that a data structure containing a generator is replaced by PicklePlaceholder rather than raising an error. """ @@ -178,13 +191,12 @@ def simple_generator(): next(reloaded["generator"]) # Attempting to call methods on the generator => AttributeError - with pytest.raises(AttributeError): + with pytest.raises(PicklePlaceholderAccessError): reloaded["generator"].send(None) def test_picklepatch_loads_standard_pickle(): - """ - Test that PicklePatcher.loads can correctly load data that was pickled + """Test that PicklePatcher.loads can correctly load data that was pickled using the standard pickle module. """ # Create a simple data structure @@ -209,12 +221,10 @@ def test_picklepatch_loads_standard_pickle(): def test_picklepatch_loads_dill_pickle(): - """ - Test that PicklePatcher.loads can correctly load data that was pickled + """Test that PicklePatcher.loads can correctly load data that was pickled using the dill module, which can pickle more complex objects than the standard pickle module. """ - # Create a more complex data structure that includes a lambda function # which dill can handle but standard pickle cannot original_data = { @@ -240,80 +250,264 @@ def test_picklepatch_loads_dill_pickle(): assert reloaded["nested"]["another_function"](4) == 16 def test_run_and_parse_picklepatch() -> None: + """Test the end to end functionality of picklepatch, from tracing benchmarks to running the replay tests. + + The first example has an argument (an object containing a socket) that is not pickleable However, the socket attributs is not used, so we are able to compare the test results with the optimized test results. + Here, we are simply 'ignoring' the unused unpickleable object. - test_path = ( - Path(__file__).parent.resolve() - / "../code_to_optimize/tests/pytest/test_bubble_sort_picklepatch.py" - ).resolve() - test_path_perf = ( - Path(__file__).parent.resolve() - / "../code_to_optimize/tests/pytest/test_bubble_sort_picklepatch_perf.py" - ).resolve() - fto_path = (Path(__file__).parent.resolve() / "../code_to_optimize/bubble_sort_picklepatch.py").resolve() - original_test =test_path.read_text("utf-8") + The second example also has an argument (an object containing socket) that is not pickleable. The socket attribute is used, which results in an error thrown by the PicklePlaceholder object. + Both the original and optimized results should error out in this case, but this should be flagged as incorrect behavior when comparing test results, + since we were not able to reuse the unpickleable object in the replay test. + """ + # Init paths + project_root = Path(__file__).parent.parent.resolve() + tests_root = project_root / "code_to_optimize" / "tests" / "pytest" + benchmarks_root = project_root / "code_to_optimize" / "tests" / "pytest" / "benchmarks_socket_test" + replay_tests_dir = benchmarks_root / "codeflash_replay_tests" + output_file = (benchmarks_root / Path("test_trace_benchmarks.trace")).resolve() + fto_unused_socket_path = (project_root / "code_to_optimize" / "bubble_sort_picklepatch_test_unused_socket.py").resolve() + fto_used_socket_path = (project_root / "code_to_optimize" / "bubble_sort_picklepatch_test_used_socket.py").resolve() + original_fto_unused_socket_code = fto_unused_socket_path.read_text("utf-8") + original_fto_used_socket_code = fto_used_socket_path.read_text("utf-8") + # Trace benchmarks + trace_benchmarks_pytest(benchmarks_root, tests_root, project_root, output_file) + assert output_file.exists() try: - tests_root = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/").resolve() - project_root_path = (Path(__file__).parent / "..").resolve() + # Check contents + conn = sqlite3.connect(output_file.as_posix()) + cursor = conn.cursor() + + 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_picklepatch_test_unused_socket.bubble_sort_with_unused_socket" in function_to_results + + test_name, total_time, function_time, percent = function_to_results["code_to_optimize.bubble_sort_picklepatch_test_unused_socket.bubble_sort_with_unused_socket"][0] + assert total_time > 0.0 + assert function_time > 0.0 + assert percent > 0.0 + + test_name, total_time, function_time, percent = \ + function_to_results["code_to_optimize.bubble_sort_picklepatch_test_unused_socket.bubble_sort_with_unused_socket"][0] + assert total_time > 0.0 + assert function_time > 0.0 + assert percent > 0.0 + + bubble_sort_unused_socket_path = (project_root / "code_to_optimize"/ "bubble_sort_picklepatch_test_unused_socket.py").as_posix() + bubble_sort_used_socket_path = (project_root / "code_to_optimize" / "bubble_sort_picklepatch_test_used_socket.py").as_posix() + # Expected function calls + expected_calls = [ + ("bubble_sort_with_unused_socket", "", "code_to_optimize.bubble_sort_picklepatch_test_unused_socket", + f"{bubble_sort_unused_socket_path}", + "test_socket_picklepatch", "code_to_optimize.tests.pytest.benchmarks_socket_test.test_socket", 12), + ("bubble_sort_with_used_socket", "", "code_to_optimize.bubble_sort_picklepatch_test_used_socket", + f"{bubble_sort_used_socket_path}", + "test_used_socket_picklepatch", "code_to_optimize.tests.pytest.benchmarks_socket_test.test_socket", 20) + ] + 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" + conn.close() + + # Generate replay test + generate_replay_test(output_file, replay_tests_dir) + replay_test_path = replay_tests_dir / Path( + "test_code_to_optimize_tests_pytest_benchmarks_socket_test_test_socket__replay_test_0.py") + replay_test_perf_path = replay_tests_dir / Path( + "test_code_to_optimize_tests_pytest_benchmarks_socket_test_test_socket__replay_test_0_perf.py") + assert replay_test_path.exists() + original_replay_test_code = replay_test_path.read_text("utf-8") + + # Instrument the replay test + func = FunctionToOptimize(function_name="bubble_sort_with_unused_socket", parents=[], file_path=Path(fto_unused_socket_path)) original_cwd = Path.cwd() - run_cwd = Path(__file__).parent.parent.resolve() - func = FunctionToOptimize(function_name="bubble_sort_with_unused_socket", parents=[], file_path=Path(fto_path)) + run_cwd = project_root os.chdir(run_cwd) success, new_test = inject_profiling_into_existing_test( - test_path, - [CodePosition(13,14), CodePosition(31,14)], + replay_test_path, + [CodePosition(17, 15)], func, - project_root_path, + project_root, "pytest", mode=TestingMode.BEHAVIOR, ) os.chdir(original_cwd) assert success assert new_test is not None - - with test_path.open("w") as f: - f.write(new_test) + replay_test_path.write_text(new_test) opt = Optimizer( Namespace( - project_root=project_root_path, + project_root=project_root, disable_telemetry=True, tests_root=tests_root, test_framework="pytest", pytest_cmd="pytest", experiment_id=None, - test_project_root=project_root_path, + test_project_root=project_root, ) ) + + # Run the replay test for the original code that does not use the socket test_env = os.environ.copy() test_env["CODEFLASH_TEST_ITERATION"] = "0" test_env["CODEFLASH_LOOP_INDEX"] = "1" - test_type = TestType.EXISTING_UNIT_TEST + test_type = TestType.REPLAY_TEST + replay_test_function = "test_code_to_optimize_bubble_sort_picklepatch_test_unused_socket_bubble_sort_with_unused_socket" + func_optimizer = opt.create_function_optimizer(func) + func_optimizer.test_files = TestFiles( + test_files=[ + TestFile( + instrumented_behavior_file_path=replay_test_path, + test_type=test_type, + original_file_path=replay_test_path, + benchmarking_file_path=replay_test_perf_path, + tests_in_file=[TestsInFile(test_file=replay_test_path, test_class=None, test_function=replay_test_function, test_type=test_type)], + ) + ] + ) + test_results_unused_socket, coverage_data = func_optimizer.run_and_parse_tests( + testing_type=TestingMode.BEHAVIOR, + test_env=test_env, + test_files=func_optimizer.test_files, + optimization_iteration=0, + pytest_min_loops=1, + pytest_max_loops=1, + testing_time=1.0, + ) + assert len(test_results_unused_socket) == 1 + assert test_results_unused_socket.test_results[0].id.test_module_path == "code_to_optimize.tests.pytest.benchmarks_socket_test.codeflash_replay_tests.test_code_to_optimize_tests_pytest_benchmarks_socket_test_test_socket__replay_test_0" + assert test_results_unused_socket.test_results[0].id.test_function_name == "test_code_to_optimize_bubble_sort_picklepatch_test_unused_socket_bubble_sort_with_unused_socket" + assert test_results_unused_socket.test_results[0].did_pass == True + + # Replace with optimized candidate + fto_unused_socket_path.write_text(""" +from codeflash.benchmarking.codeflash_trace import codeflash_trace + +@codeflash_trace +def bubble_sort_with_unused_socket(data_container): + # Extract the list to sort, leaving the socket untouched + numbers = data_container.get('numbers', []).copy() + return sorted(numbers) +""") + # Run optimized code for unused socket + optimized_test_results_unused_socket, coverage_data = func_optimizer.run_and_parse_tests( + testing_type=TestingMode.BEHAVIOR, + test_env=test_env, + test_files=func_optimizer.test_files, + optimization_iteration=0, + pytest_min_loops=1, + pytest_max_loops=1, + testing_time=1.0, + ) + assert len(optimized_test_results_unused_socket) == 1 + verification_result = compare_test_results(test_results_unused_socket, optimized_test_results_unused_socket) + assert verification_result is True + + # Remove the previous instrumentation + replay_test_path.write_text(original_replay_test_code) + # Instrument the replay test + func = FunctionToOptimize(function_name="bubble_sort_with_used_socket", parents=[], file_path=Path(fto_used_socket_path)) + success, new_test = inject_profiling_into_existing_test( + replay_test_path, + [CodePosition(23,15)], + func, + project_root, + "pytest", + mode=TestingMode.BEHAVIOR, + ) + os.chdir(original_cwd) + assert success + assert new_test is not None + replay_test_path.write_text(new_test) + # Run test for original function code that uses the socket. This should fail, as the PicklePlaceholder is accessed. + test_env = os.environ.copy() + test_env["CODEFLASH_TEST_ITERATION"] = "0" + test_env["CODEFLASH_LOOP_INDEX"] = "1" + test_type = TestType.REPLAY_TEST + func = FunctionToOptimize(function_name="bubble_sort_with_used_socket", parents=[], + file_path=Path(fto_used_socket_path)) + replay_test_function = "test_code_to_optimize_bubble_sort_picklepatch_test_used_socket_bubble_sort_with_used_socket" func_optimizer = opt.create_function_optimizer(func) func_optimizer.test_files = TestFiles( test_files=[ TestFile( - instrumented_behavior_file_path=test_path, + instrumented_behavior_file_path=replay_test_path, test_type=test_type, - original_file_path=test_path, - benchmarking_file_path=test_path_perf, + original_file_path=replay_test_path, + benchmarking_file_path=replay_test_perf_path, + tests_in_file=[ + TestsInFile(test_file=replay_test_path, test_class=None, test_function=replay_test_function, + test_type=test_type)], ) ] ) - test_results, coverage_data = func_optimizer.run_and_parse_tests( + test_results_used_socket, coverage_data = func_optimizer.run_and_parse_tests( + testing_type=TestingMode.BEHAVIOR, + test_env=test_env, + test_files=func_optimizer.test_files, + optimization_iteration=0, + pytest_min_loops=1, + pytest_max_loops=1, + testing_time=1.0, + ) + assert len(test_results_used_socket) == 1 + assert test_results_used_socket.test_results[ + 0].id.test_module_path == "code_to_optimize.tests.pytest.benchmarks_socket_test.codeflash_replay_tests.test_code_to_optimize_tests_pytest_benchmarks_socket_test_test_socket__replay_test_0" + assert test_results_used_socket.test_results[ + 0].id.test_function_name == "test_code_to_optimize_bubble_sort_picklepatch_test_used_socket_bubble_sort_with_used_socket" + assert test_results_used_socket.test_results[0].did_pass is False + print("test results used socket") + print(test_results_used_socket) + # Replace with optimized candidate + fto_used_socket_path.write_text(""" +from codeflash.benchmarking.codeflash_trace import codeflash_trace + +@codeflash_trace +def bubble_sort_with_used_socket(data_container): + # Extract the list to sort, leaving the socket untouched + numbers = data_container.get('numbers', []).copy() + socket = data_container.get('socket') + socket.send("Hello from the optimized function!") + return sorted(numbers) + """) + + # Run test for optimized function code that uses the socket. This should fail, as the PicklePlaceholder is accessed. + optimized_test_results_used_socket, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, - testing_time=0.1, + testing_time=1.0, ) - assert test_results.test_results[0].id.test_function_name =="test_bubble_sort_with_unused_socket" - assert test_results.test_results[0].did_pass ==True - assert test_results.test_results[1].id.test_function_name =="test_bubble_sort_with_used_socket" - assert test_results.test_results[1].did_pass ==False - # assert pickle placeholder problem - print(test_results) + assert len(test_results_used_socket) == 1 + assert test_results_used_socket.test_results[ + 0].id.test_module_path == "code_to_optimize.tests.pytest.benchmarks_socket_test.codeflash_replay_tests.test_code_to_optimize_tests_pytest_benchmarks_socket_test_test_socket__replay_test_0" + assert test_results_used_socket.test_results[ + 0].id.test_function_name == "test_code_to_optimize_bubble_sort_picklepatch_test_used_socket_bubble_sort_with_used_socket" + assert test_results_used_socket.test_results[0].did_pass is False + + # Even though tests threw the same error, we reject this as the behavior of the unpickleable object could not be determined. + assert compare_test_results(test_results_used_socket, optimized_test_results_used_socket) is False + finally: - test_path.write_text(original_test) \ No newline at end of file + # cleanup + output_file.unlink(missing_ok=True) + shutil.rmtree(replay_tests_dir, ignore_errors=True) + fto_unused_socket_path.write_text(original_fto_unused_socket_code) + fto_used_socket_path.write_text(original_fto_used_socket_code) + diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 715955063..af9a1e3f3 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,15 +1,14 @@ +import shutil import sqlite3 +from pathlib import Path from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin -from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest from codeflash.benchmarking.replay_test import generate_replay_test -from pathlib import Path - -from codeflash.benchmarking.utils import print_benchmark_table, validate_and_format_benchmark_table -import shutil +from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest +from codeflash.benchmarking.utils import validate_and_format_benchmark_table -def test_trace_benchmarks(): +def test_trace_benchmarks() -> None: # Test the trace_benchmarks function project_root = Path(__file__).parent.parent / "code_to_optimize" benchmarks_root = project_root / "tests" / "pytest" / "benchmarks_test" @@ -83,13 +82,12 @@ def test_trace_benchmarks(): test_class_sort_path = replay_tests_dir/ Path("test_tests_pytest_benchmarks_test_test_benchmark_bubble_sort_example__replay_test_0.py") assert test_class_sort_path.exists() test_class_sort_code = f""" -import dill as pickle - from code_to_optimize.bubble_sort_codeflash_trace import \\ Sorter as code_to_optimize_bubble_sort_codeflash_trace_Sorter from code_to_optimize.bubble_sort_codeflash_trace import \\ sorter as code_to_optimize_bubble_sort_codeflash_trace_sorter from codeflash.benchmarking.replay_test import get_next_arg_and_return +from codeflash.picklepatch.pickle_patcher import PicklePatcher as pickle functions = ['sort_class', 'sort_static', 'sorter'] trace_file_path = r"{output_file.as_posix()}" @@ -146,14 +144,13 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_Sorter___init__(): test_sort_path = replay_tests_dir / Path("test_tests_pytest_benchmarks_test_test_process_and_sort_example__replay_test_0.py") assert test_sort_path.exists() test_sort_code = f""" -import dill as pickle - from code_to_optimize.bubble_sort_codeflash_trace import \\ sorter as code_to_optimize_bubble_sort_codeflash_trace_sorter from code_to_optimize.process_and_bubble_sort_codeflash_trace import \\ compute_and_sort as \\ code_to_optimize_process_and_bubble_sort_codeflash_trace_compute_and_sort from codeflash.benchmarking.replay_test import get_next_arg_and_return +from codeflash.picklepatch.pickle_patcher import PicklePatcher as pickle functions = ['compute_and_sort', 'sorter'] trace_file_path = r"{output_file}" @@ -284,4 +281,4 @@ def test_trace_benchmark_decorator() -> None: finally: # cleanup - output_file.unlink(missing_ok=True) \ No newline at end of file + output_file.unlink(missing_ok=True) From 4bb0aadd9756237ac08e0df446248ae32510c469 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 11 Apr 2025 13:39:59 -0400 Subject: [PATCH 106/122] minor fix for removing files --- codeflash/optimization/optimizer.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index a1260dfd8..2215552a3 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -127,8 +127,6 @@ def run(self) -> None: logger.info("Information on existing benchmarks will not be available for this run.") finally: # Restore original source code - trace_file.unlink() - shutil.rmtree(self.replay_tests_dir, ignore_errors=True) for file in file_path_to_source_code: with file.open("w", encoding="utf8") as f: f.write(file_path_to_source_code[file]) From 790d77c63abf912217cfaa88a9977fd298c661eb Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Fri, 11 Apr 2025 15:49:31 -0400 Subject: [PATCH 107/122] fixes to sync with main --- .../discovery/pytest_new_process_discovery.py | 8 +- codeflash/tracer.py | 4 + codeflash/verification/test_results.py | 276 ------------------ codeflash/verification/test_runner.py | 2 - codeflash/verification/verification_utils.py | 2 - 5 files changed, 11 insertions(+), 281 deletions(-) delete mode 100644 codeflash/verification/test_results.py diff --git a/codeflash/discovery/pytest_new_process_discovery.py b/codeflash/discovery/pytest_new_process_discovery.py index d5a80f501..c16d524d4 100644 --- a/codeflash/discovery/pytest_new_process_discovery.py +++ b/codeflash/discovery/pytest_new_process_discovery.py @@ -16,6 +16,12 @@ def pytest_collection_finish(self, session) -> None: collected_tests.extend(session.items) pytest_rootdir = session.config.rootdir + def pytest_collection_modifyitems(config, items): + skip_benchmark = pytest.mark.skip(reason="Skipping benchmark tests") + for item in items: + if "benchmark" in item.fixturenames: + item.add_marker(skip_benchmark) + def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, str]]: test_results = [] @@ -34,7 +40,7 @@ def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, s try: exitcode = pytest.main( - [tests_root, "-pno:logging", "--collect-only", "-m", "not skip", "--benchmark-skip"], plugins=[PytestCollectionPlugin()] + [tests_root, "-p no:logging", "--collect-only", "-m", "not skip", "--benchmark-skip"], plugins=[PytestCollectionPlugin()] ) except Exception as e: # noqa: BLE001 print(f"Failed to collect tests: {e!s}") # noqa: T201 diff --git a/codeflash/tracer.py b/codeflash/tracer.py index eb4df84d4..5d1240868 100644 --- a/codeflash/tracer.py +++ b/codeflash/tracer.py @@ -247,14 +247,18 @@ def tracer_logic(self, frame: FrameType, event: str) -> None: return if self.timeout is not None and (time.time() - self.start_time) > self.timeout: sys.setprofile(None) + threading.setprofile(None) console.print(f"Codeflash: Timeout reached! Stopping tracing at {self.timeout} seconds.") return code = frame.f_code + file_name = Path(code.co_filename).resolve() # TODO : It currently doesn't log the last return call from the first function if code.co_name in self.ignored_functions: return + if not file_name.is_relative_to(self.project_root): + return if not file_name.exists(): return if self.functions and code.co_name not in self.functions: diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py deleted file mode 100644 index 25f258e26..000000000 --- a/codeflash/verification/test_results.py +++ /dev/null @@ -1,276 +0,0 @@ -from __future__ import annotations - -import sys -from collections import defaultdict -from enum import Enum -from pathlib import Path -from typing import TYPE_CHECKING, Optional, cast - -from pydantic import BaseModel -from pydantic.dataclasses import dataclass -from rich.tree import Tree - -from codeflash.cli_cmds.console import DEBUG_MODE, logger -from codeflash.verification.comparator import comparator - -if TYPE_CHECKING: - from collections.abc import Iterator - - -class VerificationType(str, Enum): - FUNCTION_CALL = ( - "function_call" # Correctness verification for a test function, checks input values and output values) - ) - INIT_STATE_FTO = "init_state_fto" # Correctness verification for fto class instance attributes after init - INIT_STATE_HELPER = "init_state_helper" # Correctness verification for helper class instance attributes after init - - -class TestType(Enum): - EXISTING_UNIT_TEST = 1 - INSPIRED_REGRESSION = 2 - GENERATED_REGRESSION = 3 - REPLAY_TEST = 4 - CONCOLIC_COVERAGE_TEST = 5 - INIT_STATE_TEST = 6 - - def to_name(self) -> str: - if self is TestType.INIT_STATE_TEST: - return "" - names = { - TestType.EXISTING_UNIT_TEST: "βš™οΈ Existing Unit Tests", - TestType.INSPIRED_REGRESSION: "🎨 Inspired Regression Tests", - TestType.GENERATED_REGRESSION: "πŸŒ€ Generated Regression Tests", - TestType.REPLAY_TEST: "βͺ Replay Tests", - TestType.CONCOLIC_COVERAGE_TEST: "πŸ”Ž Concolic Coverage Tests", - } - return names[self] - - -@dataclass(frozen=True) -class InvocationId: - test_module_path: str # The fully qualified name of the test module - test_class_name: Optional[str] # The name of the class where the test is defined - test_function_name: Optional[str] # The name of the test_function. Does not include the components of the file_name - function_getting_tested: str - iteration_id: Optional[str] - - # test_module_path:TestSuiteClass.test_function_name:function_tested:iteration_id - def id(self) -> str: - class_prefix = f"{self.test_class_name}." if self.test_class_name else "" - return ( - f"{self.test_module_path}:{class_prefix}{self.test_function_name}:" - f"{self.function_getting_tested}:{self.iteration_id}" - ) - - @staticmethod - def from_str_id(string_id: str, iteration_id: Optional[str] = None) -> InvocationId: - components = string_id.split(":") - assert len(components) == 4 - second_components = components[1].split(".") - if len(second_components) == 1: - test_class_name = None - test_function_name = second_components[0] - else: - test_class_name = second_components[0] - test_function_name = second_components[1] - return InvocationId( - test_module_path=components[0], - test_class_name=test_class_name, - test_function_name=test_function_name, - function_getting_tested=components[2], - iteration_id=iteration_id if iteration_id else components[3], - ) - - -@dataclass(frozen=True) -class FunctionTestInvocation: - loop_index: int # The loop index of the function invocation, starts at 1 - id: InvocationId # The fully qualified name of the function invocation (id) - file_name: Path # The file where the test is defined - did_pass: bool # Whether the test this function invocation was part of, passed or failed - runtime: Optional[int] # Time in nanoseconds - test_framework: str # unittest or pytest - test_type: TestType - return_value: Optional[object] # The return value of the function invocation - timed_out: Optional[bool] - verification_type: Optional[str] = VerificationType.FUNCTION_CALL - stdout: Optional[str] = None - - @property - def unique_invocation_loop_id(self) -> str: - return f"{self.loop_index}:{self.id.id()}" - - -class TestResults(BaseModel): - # don't modify these directly, use the add method - # also we don't support deletion of test results elements - caution is advised - test_results: list[FunctionTestInvocation] = [] - test_result_idx: dict[str, int] = {} - - def add(self, function_test_invocation: FunctionTestInvocation) -> None: - unique_id = function_test_invocation.unique_invocation_loop_id - if unique_id in self.test_result_idx: - if DEBUG_MODE: - logger.warning(f"Test result with id {unique_id} already exists. SKIPPING") - return - self.test_result_idx[unique_id] = len(self.test_results) - self.test_results.append(function_test_invocation) - - def merge(self, other: TestResults) -> None: - original_len = len(self.test_results) - self.test_results.extend(other.test_results) - for k, v in other.test_result_idx.items(): - if k in self.test_result_idx: - msg = f"Test result with id {k} already exists." - raise ValueError(msg) - self.test_result_idx[k] = v + original_len - - def filter_by_test_type(self, test_type: TestType) -> TestResults: - filtered_test_results = [] - filtered_test_results_idx = {} - for test_result in self.test_results: - if test_result.test_type == test_type: - filtered_test_results_idx[test_result.unique_invocation_loop_id] = len(filtered_test_results) - filtered_test_results.append(test_result) - return TestResults(test_results=filtered_test_results, test_result_idx=filtered_test_results_idx) - - def group_by_benchmark(self, benchmark_key_set:set[tuple[str,str]]) -> dict[tuple[str,str],TestResults]: - """Group TestResults by benchmark key. - - For now, use a tuple of (file_path, function_name) as the benchmark key. Can't import BenchmarkKey because of circular import. - - Args: - benchmark_key_set (set[tuple[str,str]]): A set of tuples of (file_path, function_name) - - Returns: - TestResults: A new TestResults object with the test results grouped by benchmark key. - - """ - test_result_by_benchmark = defaultdict(TestResults) - for test_result in self.test_results: - if test_result.test_type == TestType.REPLAY_TEST and (test_result.id.test_module_path,test_result.id.test_function_name) in benchmark_key_set: - test_result_by_benchmark[(test_result.id.test_module_path,test_result.id.test_function_name)].add(test_result) - return test_result_by_benchmark - - def get_by_unique_invocation_loop_id(self, unique_invocation_loop_id: str) -> FunctionTestInvocation | None: - try: - return self.test_results[self.test_result_idx[unique_invocation_loop_id]] - except (IndexError, KeyError): - return None - - def get_all_ids(self) -> set[InvocationId]: - return {test_result.id for test_result in self.test_results} - - def get_all_unique_invocation_loop_ids(self) -> set[str]: - return {test_result.unique_invocation_loop_id for test_result in self.test_results} - - def number_of_loops(self) -> int: - if not self.test_results: - return 0 - return max(test_result.loop_index for test_result in self.test_results) - - def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]: - report = {} - for test_type in TestType: - report[test_type] = {"passed": 0, "failed": 0} - for test_result in self.test_results: - if test_result.loop_index == 1: - if test_result.did_pass: - report[test_result.test_type]["passed"] += 1 - else: - report[test_result.test_type]["failed"] += 1 - return report - - @staticmethod - def report_to_string(report: dict[TestType, dict[str, int]]) -> str: - return " ".join( - [ - f"{test_type.to_name()}- (Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']})" - for test_type in TestType - ] - ) - - @staticmethod - def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree: - tree = Tree(title) - for test_type in TestType: - if test_type is TestType.INIT_STATE_TEST: - continue - tree.add( - f"{test_type.to_name()} - Passed: {report[test_type]['passed']}, Failed: {report[test_type]['failed']}" - ) - return tree - - def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: - usable_runtime_by_id = defaultdict(list) - for result in self.test_results: - if result.did_pass: - if not result.runtime: - msg = ( - f"Ignoring test case that passed but had no runtime -> {result.id}, " - f"Loop # {result.loop_index}, Test Type: {result.test_type}, " - f"Verification Type: {result.verification_type}" - ) - logger.debug(msg) - else: - usable_runtime_by_id[result.id].append(result.runtime) - - return usable_runtime_by_id - - - def total_passed_runtime(self) -> int: - """Calculate the sum of runtimes of all test cases that passed. - - A testcase runtime is the minimum value of all looped execution runtimes. - - :return: The runtime in nanoseconds. - """ - return sum( - [min(usable_runtime_data) for _, usable_runtime_data in self.usable_runtime_data_by_test_case().items()] - ) - - def __iter__(self) -> Iterator[FunctionTestInvocation]: - return iter(self.test_results) - - def __len__(self) -> int: - return len(self.test_results) - - def __getitem__(self, index: int) -> FunctionTestInvocation: - return self.test_results[index] - - def __setitem__(self, index: int, value: FunctionTestInvocation) -> None: - self.test_results[index] = value - - def __contains__(self, value: FunctionTestInvocation) -> bool: - return value in self.test_results - - def __bool__(self) -> bool: - return bool(self.test_results) - - def __eq__(self, other: object) -> bool: - # Unordered comparison - if type(self) is not type(other): - return False - if len(self) != len(other): - return False - original_recursion_limit = sys.getrecursionlimit() - cast(TestResults, other) - for test_result in self: - other_test_result = other.get_by_unique_invocation_loop_id(test_result.unique_invocation_loop_id) - if other_test_result is None: - return False - - if original_recursion_limit < 5000: - sys.setrecursionlimit(5000) - if ( - test_result.file_name != other_test_result.file_name - or test_result.did_pass != other_test_result.did_pass - or test_result.runtime != other_test_result.runtime - or test_result.test_framework != other_test_result.test_framework - or test_result.test_type != other_test_result.test_type - or not comparator(test_result.return_value, other_test_result.return_value) - ): - sys.setrecursionlimit(original_recursion_limit) - return False - sys.setrecursionlimit(original_recursion_limit) - return True diff --git a/codeflash/verification/test_runner.py b/codeflash/verification/test_runner.py index 3d4780e81..483695b1a 100644 --- a/codeflash/verification/test_runner.py +++ b/codeflash/verification/test_runner.py @@ -62,8 +62,6 @@ def run_behavioral_tests( "--capture=tee-sys", f"--timeout={pytest_timeout}", "-q", - "-o", - "addopts=", "--codeflash_loops_scope=session", "--codeflash_min_loops=1", "--codeflash_max_loops=1", diff --git a/codeflash/verification/verification_utils.py b/codeflash/verification/verification_utils.py index 53a756f27..43cb78770 100644 --- a/codeflash/verification/verification_utils.py +++ b/codeflash/verification/verification_utils.py @@ -6,8 +6,6 @@ from pydantic.dataclasses import dataclass -from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE - def get_test_file_path(test_dir: Path, function_name: str, iteration: int = 0, test_type: str = "unit") -> Path: assert test_type in {"unit", "inspired", "replay", "perf"} From 28fd746d01b33f3ea5a28bfbbbeb47bbeeb254ac Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 15 Apr 2025 13:26:35 -0400 Subject: [PATCH 108/122] cmd init changes --- codeflash/cli_cmds/cmd_init.py | 85 +++++++++++++++++++++++++++++++--- codeflash/models/models.py | 28 +++++------ 2 files changed, 94 insertions(+), 19 deletions(-) diff --git a/codeflash/cli_cmds/cmd_init.py b/codeflash/cli_cmds/cmd_init.py index 6c7615fcd..6e1b2c240 100644 --- a/codeflash/cli_cmds/cmd_init.py +++ b/codeflash/cli_cmds/cmd_init.py @@ -50,6 +50,7 @@ class SetupInfo: module_root: str tests_root: str + benchmarks_root: str | None test_framework: str ignore_paths: list[str] formatter: str @@ -126,8 +127,7 @@ def ask_run_end_to_end_test(args: Namespace) -> None: run_end_to_end_test(args, bubble_sort_path, bubble_sort_test_path) def should_modify_pyproject_toml() -> bool: - """ - Check if the current directory contains a valid pyproject.toml file with codeflash config + """Check if the current directory contains a valid pyproject.toml file with codeflash config If it does, ask the user if they want to re-configure it. """ from rich.prompt import Confirm @@ -136,7 +136,7 @@ def should_modify_pyproject_toml() -> bool: return True try: config, config_file_path = parse_config_file(pyproject_toml_path) - except Exception as e: + except Exception: return True if "module_root" not in config or config["module_root"] is None or not Path(config["module_root"]).is_dir(): @@ -145,7 +145,7 @@ def should_modify_pyproject_toml() -> bool: return True create_toml = Confirm.ask( - f"βœ… A valid Codeflash config already exists in this project. Do you want to re-configure it?", default=False, show_default=True + "βœ… A valid Codeflash config already exists in this project. Do you want to re-configure it?", default=False, show_default=True ) return create_toml @@ -245,6 +245,66 @@ def collect_setup_info() -> SetupInfo: ph("cli-test-framework-provided", {"test_framework": test_framework}) + # Get benchmarks root directory + default_benchmarks_subdir = "benchmarks" + create_benchmarks_option = f"okay, create a {default_benchmarks_subdir}{os.path.sep} directory for me!" + no_benchmarks_option = "I don't need benchmarks" + + # Check if benchmarks directory exists inside tests directory + tests_subdirs = [] + if tests_root.exists(): + tests_subdirs = [d.name for d in tests_root.iterdir() if d.is_dir() and not d.name.startswith(".")] + + benchmarks_options = [] + if default_benchmarks_subdir in tests_subdirs: + benchmarks_options.append(default_benchmarks_subdir) + benchmarks_options.extend([d for d in tests_subdirs if d != default_benchmarks_subdir]) + benchmarks_options.append(create_benchmarks_option) + benchmarks_options.append(custom_dir_option) + benchmarks_options.append(no_benchmarks_option) + + benchmarks_answer = inquirer_wrapper( + inquirer.list_input, + message="Where are your benchmarks located? (benchmarks must be a sub directory of your tests root directory)", + choices=benchmarks_options, + default=( + default_benchmarks_subdir if default_benchmarks_subdir in benchmarks_options else benchmarks_options[0]), + ) + + if benchmarks_answer == create_benchmarks_option: + benchmarks_root = tests_root / default_benchmarks_subdir + benchmarks_root.mkdir(exist_ok=True) + click.echo(f"βœ… Created directory {benchmarks_root}{os.path.sep}{LF}") + elif benchmarks_answer == custom_dir_option: + custom_benchmarks_answer = inquirer_wrapper_path( + "path", + message=f"Enter the path to your benchmarks directory inside {tests_root}{os.path.sep} ", + path_type=inquirer.Path.DIRECTORY, + ) + if custom_benchmarks_answer: + benchmarks_root = tests_root / Path(custom_benchmarks_answer["path"]) + else: + apologize_and_exit() + elif benchmarks_answer == no_benchmarks_option: + benchmarks_root = None + else: + benchmarks_root = tests_root / Path(cast(str, benchmarks_answer)) + + # TODO: Implement other benchmark framework options + # if benchmarks_root: + # benchmarks_root = benchmarks_root.relative_to(curdir) + # + # # Ask about benchmark framework + # benchmark_framework_options = ["pytest-benchmark", "asv (Airspeed Velocity)", "custom/other"] + # benchmark_framework = inquirer_wrapper( + # inquirer.list_input, + # message="Which benchmark framework do you use?", + # choices=benchmark_framework_options, + # default=benchmark_framework_options[0], + # carousel=True, + # ) + + formatter = inquirer_wrapper( inquirer.list_input, message="Which code formatter do you use?", @@ -280,6 +340,7 @@ def collect_setup_info() -> SetupInfo: return SetupInfo( module_root=str(module_root), tests_root=str(tests_root), + benchmarks_root = str(benchmarks_root) if benchmarks_root else None, test_framework=cast(str, test_framework), ignore_paths=ignore_paths, formatter=cast(str, formatter), @@ -438,11 +499,19 @@ def install_github_actions(override_formatter_check: bool = False) -> None: return workflows_path.mkdir(parents=True, exist_ok=True) from importlib.resources import files + benchmark_mode = False + if "benchmarks_root" in config: + benchmark_mode = inquirer_wrapper( + inquirer.confirm, + message="⚑️It looks like you've configured a benchmarks_root in your config. Would you like to run the Github action in benchmark mode? " + " This will show the impact of Codeflash's suggested optimizations on your benchmarks", + default=True, + ) optimize_yml_content = ( files("codeflash").joinpath("cli_cmds", "workflows", "codeflash-optimize.yaml").read_text(encoding="utf-8") ) - materialized_optimize_yml_content = customize_codeflash_yaml_content(optimize_yml_content, config, git_root) + materialized_optimize_yml_content = customize_codeflash_yaml_content(optimize_yml_content, config, git_root, benchmark_mode) with optimize_yaml_path.open("w", encoding="utf8") as optimize_yml_file: optimize_yml_file.write(materialized_optimize_yml_content) click.echo(f"{LF}βœ… Created GitHub action workflow at {optimize_yaml_path}{LF}") @@ -557,7 +626,7 @@ def get_github_action_working_directory(toml_path: Path, git_root: Path) -> str: def customize_codeflash_yaml_content( - optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path + optimize_yml_content: str, config: tuple[dict[str, Any], Path], git_root: Path, benchmark_mode: bool = False ) -> str: module_path = str(Path(config["module_root"]).relative_to(git_root) / "**") optimize_yml_content = optimize_yml_content.replace("{{ codeflash_module_path }}", module_path) @@ -588,6 +657,9 @@ def customize_codeflash_yaml_content( # Add codeflash command codeflash_cmd = get_codeflash_github_action_command(dep_manager) + + if benchmark_mode: + codeflash_cmd += " --benchmark" return optimize_yml_content.replace("{{ codeflash_command }}", codeflash_cmd) @@ -609,6 +681,7 @@ def configure_pyproject_toml(setup_info: SetupInfo) -> None: codeflash_section["module-root"] = setup_info.module_root codeflash_section["tests-root"] = setup_info.tests_root codeflash_section["test-framework"] = setup_info.test_framework + codeflash_section["benchmarks-root"] = setup_info.benchmarks_root if setup_info.benchmarks_root else "" codeflash_section["ignore-paths"] = setup_info.ignore_paths if setup_info.git_remote not in ["", "origin"]: codeflash_section["git-remote"] = setup_info.git_remote diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 791912b8a..ddaccd16e 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -537,20 +537,22 @@ def report_to_tree(report: dict[TestType, dict[str, int]], title: str) -> Tree: return tree def usable_runtime_data_by_test_case(self) -> dict[InvocationId, list[int]]: - - usable_runtime = defaultdict(list) for result in self.test_results: - if result.did_pass: - if not result.runtime: - msg = ( - f"Ignoring test case that passed but had no runtime -> {result.id}, " - f"Loop # {result.loop_index}, Test Type: {result.test_type}, " - f"Verification Type: {result.verification_type}" - ) - logger.debug(msg) - else: - usable_runtime[result.id].append(result.runtime) - return usable_runtime + if result.did_pass and not result.runtime: + msg = ( + f"Ignoring test case that passed but had no runtime -> {result.id}, " + f"Loop # {result.loop_index}, Test Type: {result.test_type}, " + f"Verification Type: {result.verification_type}" + ) + logger.debug(msg) + + usable_runtimes = [ + (result.id, result.runtime) for result in self.test_results if result.did_pass and result.runtime + ] + return { + usable_id: [runtime[1] for runtime in usable_runtimes if runtime[0] == usable_id] + for usable_id in {runtime[0] for runtime in usable_runtimes} + } def total_passed_runtime(self) -> int: """Calculate the sum of runtimes of all test cases that passed. From 4e8483bef0e26846eacd2f4788f4b9502163f901 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 15 Apr 2025 21:18:59 -0400 Subject: [PATCH 109/122] created benchmarks for codeflash, modified codeflash-optimize to use codeflash --benchmark --- .github/workflows/codeflash-optimize.yaml | 2 +- ...est_benchmark_code_extract_code_context.py | 31 ++++++++ .../test_benchmark_discover_unit_tests.py | 26 +++++++ .../test_benchmark_merge_test_results.py | 71 +++++++++++++++++++ 4 files changed, 129 insertions(+), 1 deletion(-) create mode 100644 tests/benchmarks/test_benchmark_code_extract_code_context.py create mode 100644 tests/benchmarks/test_benchmark_discover_unit_tests.py create mode 100644 tests/benchmarks/test_benchmark_merge_test_results.py diff --git a/.github/workflows/codeflash-optimize.yaml b/.github/workflows/codeflash-optimize.yaml index 6a08635bf..357269116 100644 --- a/.github/workflows/codeflash-optimize.yaml +++ b/.github/workflows/codeflash-optimize.yaml @@ -68,4 +68,4 @@ jobs: id: optimize_code run: | source .venv/bin/activate - poetry run codeflash + poetry run codeflash --benchmark diff --git a/tests/benchmarks/test_benchmark_code_extract_code_context.py b/tests/benchmarks/test_benchmark_code_extract_code_context.py new file mode 100644 index 000000000..122276408 --- /dev/null +++ b/tests/benchmarks/test_benchmark_code_extract_code_context.py @@ -0,0 +1,31 @@ +from argparse import Namespace +from pathlib import Path + +from codeflash.context.code_context_extractor import get_code_optimization_context +from codeflash.discovery.functions_to_optimize import FunctionToOptimize +from codeflash.models.models import FunctionParent +from codeflash.optimization.optimizer import Optimizer + + +def test_benchmark_extract(benchmark)->None: + file_path = Path(__file__).parent.parent.parent.resolve() / "codeflash" + opt = Optimizer( + Namespace( + project_root=file_path.resolve(), + disable_telemetry=True, + tests_root=(file_path / "tests").resolve(), + test_framework="pytest", + pytest_cmd="pytest", + experiment_id=None, + test_project_root=Path.cwd(), + ) + ) + function_to_optimize = FunctionToOptimize( + function_name="replace_function_and_helpers_with_optimized_code", + file_path=file_path / "optimization" / "function_optimizer.py", + parents=[FunctionParent(name="FunctionOptimizer", type="ClassDef")], + starting_line=None, + ending_line=None, + ) + + benchmark(get_code_optimization_context,function_to_optimize, opt.args.project_root) diff --git a/tests/benchmarks/test_benchmark_discover_unit_tests.py b/tests/benchmarks/test_benchmark_discover_unit_tests.py new file mode 100644 index 000000000..4b05f663b --- /dev/null +++ b/tests/benchmarks/test_benchmark_discover_unit_tests.py @@ -0,0 +1,26 @@ +from pathlib import Path + +from codeflash.discovery.discover_unit_tests import discover_unit_tests +from codeflash.verification.verification_utils import TestConfig + + +def test_benchmark_code_to_optimize_test_discovery(benchmark) -> None: + project_path = Path(__file__).parent.parent.parent.resolve() / "code_to_optimize" + tests_path = project_path / "tests" / "pytest" + test_config = TestConfig( + tests_root=tests_path, + project_root_path=project_path, + test_framework="pytest", + tests_project_rootdir=tests_path.parent, + ) + benchmark(discover_unit_tests, test_config) +def test_benchmark_codeflash_test_discovery(benchmark) -> None: + project_path = Path(__file__).parent.parent.parent.resolve() / "codeflash" + tests_path = project_path / "tests" + test_config = TestConfig( + tests_root=tests_path, + project_root_path=project_path, + test_framework="pytest", + tests_project_rootdir=tests_path.parent, + ) + benchmark(discover_unit_tests, test_config) diff --git a/tests/benchmarks/test_benchmark_merge_test_results.py b/tests/benchmarks/test_benchmark_merge_test_results.py new file mode 100644 index 000000000..f0c126f75 --- /dev/null +++ b/tests/benchmarks/test_benchmark_merge_test_results.py @@ -0,0 +1,71 @@ +from codeflash.models.models import FunctionTestInvocation, InvocationId, TestResults, TestType +from codeflash.verification.parse_test_output import merge_test_results + + +def generate_test_invocations(count=100): + """Generate a set number of test invocations for benchmarking.""" + test_results_xml = TestResults() + test_results_bin = TestResults() + + # Generate test invocations in a loop + for i in range(count): + iteration_id = str(i * 3 + 5) # Generate unique iteration IDs + + # XML results - some with None runtime + test_results_xml.add( + FunctionTestInvocation( + id=InvocationId( + test_module_path="code_to_optimize.tests.unittest.test_bubble_sort", + test_class_name="TestPigLatin", + test_function_name="test_sort", + function_getting_tested="sorter", + iteration_id=iteration_id, + ), + file_name="/tmp/tests/unittest/test_bubble_sort__perfinstrumented.py", + did_pass=True, + runtime=None if i % 3 == 0 else i * 100, # Vary runtime values + test_framework="unittest", + test_type=TestType.EXISTING_UNIT_TEST, + return_value=None, + timed_out=False, + loop_index=i, + ) + ) + + # Binary results - with actual runtime values + test_results_bin.add( + FunctionTestInvocation( + id=InvocationId( + test_module_path="code_to_optimize.tests.unittest.test_bubble_sort", + test_class_name="TestPigLatin", + test_function_name="test_sort", + function_getting_tested="sorter", + iteration_id=iteration_id, + ), + file_name="/tmp/tests/unittest/test_bubble_sort__perfinstrumented.py", + did_pass=True, + runtime=500 + i * 20, # Generate varying runtime values + test_framework="unittest", + test_type=TestType.EXISTING_UNIT_TEST, + return_value=None, + timed_out=False, + loop_index=i, + ) + ) + + return test_results_xml, test_results_bin + + +def run_merge_benchmark(count=100): + test_results_xml, test_results_bin = generate_test_invocations(count) + + # Perform the merge operation that will be benchmarked + merge_test_results( + xml_test_results=test_results_xml, + bin_test_results=test_results_bin, + test_framework="unittest" + ) + + +def test_benchmark_merge_test_results(benchmark): + benchmark(run_merge_benchmark, 1000) # Default to 100 test invocations From 0680f793c963c68c45b984caade7cd4bb4cfee8b Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 15 Apr 2025 21:33:38 -0400 Subject: [PATCH 110/122] added benchmarks root --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 3ad68367a..84ec74a5e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -218,6 +218,7 @@ initial-content = """ [tool.codeflash] module-root = "codeflash" tests-root = "tests" +benchmarks-root = "tests/benchmarks" test-framework = "pytest" formatter-cmds = [ "uvx ruff check --exit-zero --fix $file", From 583b46471ade2ce3ece61e4d5e321ea94211e7f7 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Tue, 15 Apr 2025 21:46:55 -0400 Subject: [PATCH 111/122] removed comment --- codeflash/optimization/optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index 2215552a3..82e72ad97 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -82,7 +82,7 @@ def run(self) -> None: function_optimizer = None file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]] num_optimizable_functions: int - # if self.args.benchmark: + # discover functions (file_to_funcs_to_optimize, num_optimizable_functions) = get_functions_to_optimize( optimize_all=self.args.all, From 1eaaad7f973e5cfad085b65d33446ba1bfed3f89 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 16 Apr 2025 18:02:30 -0400 Subject: [PATCH 112/122] debugging --- codeflash/discovery/discover_unit_tests.py | 1 + 1 file changed, 1 insertion(+) diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py index e26680e1a..aa6725175 100644 --- a/codeflash/discovery/discover_unit_tests.py +++ b/codeflash/discovery/discover_unit_tests.py @@ -72,6 +72,7 @@ def discover_tests_pytest( capture_output=True, text=True, ) + print(result.stdout) try: with tmp_pickle_path.open(mode="rb") as f: exitcode, tests, pytest_rootdir = pickle.load(f) From ab9079b2c6f8ad049d342383a6d1d56a4642a1bc Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 16 Apr 2025 18:07:32 -0400 Subject: [PATCH 113/122] debugging --- codeflash/discovery/discover_unit_tests.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py index aa6725175..bfeb92930 100644 --- a/codeflash/discovery/discover_unit_tests.py +++ b/codeflash/discovery/discover_unit_tests.py @@ -72,7 +72,8 @@ def discover_tests_pytest( capture_output=True, text=True, ) - print(result.stdout) + logger.info(result.stdout) + logger.info(result.stderr) try: with tmp_pickle_path.open(mode="rb") as f: exitcode, tests, pytest_rootdir = pickle.load(f) From d7274ec996a4e7f0753db9dd338a12c24c2cfee4 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 16 Apr 2025 18:14:58 -0400 Subject: [PATCH 114/122] removed benchmark-skip --- codeflash/discovery/discover_unit_tests.py | 2 -- codeflash/discovery/pytest_new_process_discovery.py | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/codeflash/discovery/discover_unit_tests.py b/codeflash/discovery/discover_unit_tests.py index bfeb92930..e26680e1a 100644 --- a/codeflash/discovery/discover_unit_tests.py +++ b/codeflash/discovery/discover_unit_tests.py @@ -72,8 +72,6 @@ def discover_tests_pytest( capture_output=True, text=True, ) - logger.info(result.stdout) - logger.info(result.stderr) try: with tmp_pickle_path.open(mode="rb") as f: exitcode, tests, pytest_rootdir = pickle.load(f) diff --git a/codeflash/discovery/pytest_new_process_discovery.py b/codeflash/discovery/pytest_new_process_discovery.py index c16d524d4..2d8583255 100644 --- a/codeflash/discovery/pytest_new_process_discovery.py +++ b/codeflash/discovery/pytest_new_process_discovery.py @@ -40,7 +40,7 @@ def parse_pytest_collection_results(pytest_tests: list[Any]) -> list[dict[str, s try: exitcode = pytest.main( - [tests_root, "-p no:logging", "--collect-only", "-m", "not skip", "--benchmark-skip"], plugins=[PytestCollectionPlugin()] + [tests_root, "-p no:logging", "--collect-only", "-m", "not skip"], plugins=[PytestCollectionPlugin()] ) except Exception as e: # noqa: BLE001 print(f"Failed to collect tests: {e!s}") # noqa: T201 From a624221642b89b077cbf296d4e3b2ace27ee0153 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 16 Apr 2025 19:14:55 -0400 Subject: [PATCH 115/122] added pytest-benchmark as dependency --- .github/workflows/unit-tests.yaml | 2 +- pyproject.toml | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/unit-tests.yaml b/.github/workflows/unit-tests.yaml index a1e7da8ea..8cdcd54d9 100644 --- a/.github/workflows/unit-tests.yaml +++ b/.github/workflows/unit-tests.yaml @@ -32,7 +32,7 @@ jobs: run: uvx poetry install --with dev - name: Unit tests - run: uvx poetry run pytest tests/ --cov --cov-report=xml + run: uvx poetry run pytest tests/ --cov --cov-report=xml --benchmark-skip - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v5 diff --git a/pyproject.toml b/pyproject.toml index 84ec74a5e..d28d2469e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -115,6 +115,7 @@ types-openpyxl = ">=3.1.5.20241020" types-regex = ">=2024.9.11.20240912" types-python-dateutil = ">=2.9.0.20241003" pytest-cov = "^6.0.0" +pytest-benchmark = ">=5.1.0" types-gevent = "^24.11.0.20241230" types-greenlet = "^3.1.0.20241221" types-pexpect = "^4.9.0.20241208" From 605d078f9fb81f7aabe952c38de20f08cfd141aa Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 16 Apr 2025 19:17:27 -0400 Subject: [PATCH 116/122] updated pyproject --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index d28d2469e..bf1718d33 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,6 +92,7 @@ rich = ">=13.8.1" lxml = ">=5.3.0" crosshair-tool = ">=0.0.78" coverage = ">=7.6.4" +line_profiler=">=4.2.0" #this is the minimum version which supports python 3.13 [tool.poetry.group.dev] optional = true @@ -152,7 +153,7 @@ warn_required_dynamic_aliases = true line-length = 120 fix = true show-fixes = true -exclude = ["code_to_optimize/", "pie_test_set/"] +exclude = ["code_to_optimize/", "pie_test_set/", "tests/"] [tool.ruff.lint] select = ["ALL"] From 78871fe1dfb6b4f1a6fb45a0bff1bc49ba572686 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Wed, 16 Apr 2025 19:27:45 -0400 Subject: [PATCH 117/122] gha failing on multithreaded t est --- code_to_optimize/bubble_sort_multithread.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/code_to_optimize/bubble_sort_multithread.py b/code_to_optimize/bubble_sort_multithread.py index 3659b01bf..e71be4816 100644 --- a/code_to_optimize/bubble_sort_multithread.py +++ b/code_to_optimize/bubble_sort_multithread.py @@ -8,7 +8,7 @@ def multithreaded_sorter(unsorted_lists: list[list[int]]) -> list[list[int]]: sorted_lists = [None] * len(unsorted_lists) # Use ThreadPoolExecutor to manage threads - with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor: + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor: # Submit all sorting tasks and map them to their original indices future_to_index = { executor.submit(sorter, unsorted_list): i From 0146d828bc71419c584447fb63b860ab42c67d0b Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 17 Apr 2025 10:16:56 -0400 Subject: [PATCH 118/122] line number test is off by 1 for python versions 39 and 310, removed the check --- tests/test_trace_benchmarks.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index af9a1e3f3..a69bfabf0 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -275,7 +275,6 @@ def test_trace_benchmark_decorator() -> None: 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() From 3017ccf33a708628f63e842288a6c966524317ac Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 17 Apr 2025 10:44:41 -0400 Subject: [PATCH 119/122] 100 max function calls before flushing to disk instead of 1000 --- codeflash/benchmarking/codeflash_trace.py | 4 ++-- codeflash/picklepatch/pickle_placeholder.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/codeflash/benchmarking/codeflash_trace.py b/codeflash/benchmarking/codeflash_trace.py index 2694532f3..35232f954 100644 --- a/codeflash/benchmarking/codeflash_trace.py +++ b/codeflash/benchmarking/codeflash_trace.py @@ -160,8 +160,8 @@ def wrapper(*args, **kwargs): overhead_time, None, None) ) return result - # Flush to database every 1000 calls - if len(self.function_calls_data) > 1000: + # Flush to database every 100 calls + if len(self.function_calls_data) > 100: self.write_function_timings() # Add to the list of function calls with pickled args, to be used for replay tests diff --git a/codeflash/picklepatch/pickle_placeholder.py b/codeflash/picklepatch/pickle_placeholder.py index a422abb45..0d730dabb 100644 --- a/codeflash/picklepatch/pickle_placeholder.py +++ b/codeflash/picklepatch/pickle_placeholder.py @@ -7,7 +7,7 @@ class PicklePlaceholder: """A placeholder for an object that couldn't be pickled. When unpickled, any attempt to access attributes or call methods on this - placeholder will raise an informative exception. + placeholder will raise a PicklePlaceholderAccessError. """ def __init__(self, obj_type, obj_str, error_msg, path=None): From f14cf010eef64827a54219e25bcba2bf0e9a03c6 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 17 Apr 2025 10:56:20 -0400 Subject: [PATCH 120/122] skip multithreaded benchmark test if machine is single threaded (fixes flaky github actions test, as sometimes the machines allocated are different) --- tests/test_trace_benchmarks.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index a69bfabf0..5662a1f20 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -1,7 +1,10 @@ +import multiprocessing import shutil import sqlite3 from pathlib import Path +import pytest + from codeflash.benchmarking.plugin.plugin import codeflash_benchmark_plugin from codeflash.benchmarking.replay_test import generate_replay_test from codeflash.benchmarking.trace_benchmarks import trace_benchmarks_pytest @@ -174,6 +177,11 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): output_file.unlink(missing_ok=True) shutil.rmtree(replay_tests_dir) +# Skip the test if the machine has only 1 thread/CPU +@pytest.mark.skipif( + multiprocessing.cpu_count() <= 1, + reason="This test requires more than 1 CPU thread" +) def test_trace_multithreaded_benchmark() -> None: project_root = Path(__file__).parent.parent / "code_to_optimize" benchmarks_root = project_root / "tests" / "pytest" / "benchmarks_multithread" From e5ca10fbb660185c3927d6a13f06a36716e5329d Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 17 Apr 2025 11:11:12 -0400 Subject: [PATCH 121/122] marked multithreaded trace benchmarks test to be skipped during CI as its flaky with github action machines --- .github/workflows/unit-tests.yaml | 2 +- tests/test_trace_benchmarks.py | 7 ++----- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/.github/workflows/unit-tests.yaml b/.github/workflows/unit-tests.yaml index 8cdcd54d9..f3b4ffca5 100644 --- a/.github/workflows/unit-tests.yaml +++ b/.github/workflows/unit-tests.yaml @@ -32,7 +32,7 @@ jobs: run: uvx poetry install --with dev - name: Unit tests - run: uvx poetry run pytest tests/ --cov --cov-report=xml --benchmark-skip + run: uvx poetry run pytest tests/ --cov --cov-report=xml --benchmark-skip -m "not ci_skip" - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v5 diff --git a/tests/test_trace_benchmarks.py b/tests/test_trace_benchmarks.py index 5662a1f20..72c0267a8 100644 --- a/tests/test_trace_benchmarks.py +++ b/tests/test_trace_benchmarks.py @@ -177,11 +177,8 @@ def test_code_to_optimize_bubble_sort_codeflash_trace_sorter(): output_file.unlink(missing_ok=True) shutil.rmtree(replay_tests_dir) -# Skip the test if the machine has only 1 thread/CPU -@pytest.mark.skipif( - multiprocessing.cpu_count() <= 1, - reason="This test requires more than 1 CPU thread" -) +# Skip the test in CI as the machine may not be multithreaded +@pytest.mark.ci_skip def test_trace_multithreaded_benchmark() -> None: project_root = Path(__file__).parent.parent / "code_to_optimize" benchmarks_root = project_root / "tests" / "pytest" / "benchmarks_multithread" From 683c9f64a13153f894900629b0e97b366b019ee6 Mon Sep 17 00:00:00 2001 From: Alvin Ryanputra Date: Thu, 17 Apr 2025 16:01:33 -0400 Subject: [PATCH 122/122] shift check for pickle placerholder access error in comparator --- codeflash/verification/comparator.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/codeflash/verification/comparator.py b/codeflash/verification/comparator.py index 8a7048c57..0ebd2cc7d 100644 --- a/codeflash/verification/comparator.py +++ b/codeflash/verification/comparator.py @@ -65,11 +65,7 @@ def comparator(orig: Any, new: Any, superset_obj=False) -> bool: if len(orig) != len(new): return False return all(comparator(elem1, elem2, superset_obj) for elem1, elem2 in zip(orig, new)) - if isinstance(orig, PicklePlaceholderAccessError) or isinstance(new, PicklePlaceholderAccessError): - # If this error was raised, there was an attempt to access the PicklePlaceholder, which represents an unpickleable object. - # The test results should be rejected as the behavior of the unpickleable object is unknown. - logger.debug("Unable to verify behavior of unpickleable object in replay test") - return False + if isinstance( orig, ( @@ -95,6 +91,11 @@ def comparator(orig: Any, new: Any, superset_obj=False) -> bool: return True return math.isclose(orig, new) if isinstance(orig, BaseException): + if isinstance(orig, PicklePlaceholderAccessError) or isinstance(new, PicklePlaceholderAccessError): + # If this error was raised, there was an attempt to access the PicklePlaceholder, which represents an unpickleable object. + # The test results should be rejected as the behavior of the unpickleable object is unknown. + logger.debug("Unable to verify behavior of unpickleable object in replay test") + return False # if str(orig) != str(new): # return False # compare the attributes of the two exception objects to determine if they are equivalent.