diff --git a/codeflash/code_utils/coverage_utils.py b/codeflash/code_utils/coverage_utils.py index b115e3df8..34c5475ec 100644 --- a/codeflash/code_utils/coverage_utils.py +++ b/codeflash/code_utils/coverage_utils.py @@ -25,7 +25,18 @@ def extract_dependent_function(main_function: str, code_context: CodeOptimizatio if len(dependent_functions) != 1: return False - return dependent_functions.pop() + return build_fully_qualified_name(dependent_functions.pop(), code_context) + + +def build_fully_qualified_name(function_name: str, code_context: CodeOptimizationContext) -> str: + full_name = function_name + for obj_name, parents in code_context.preexisting_objects: + if obj_name == function_name: + for parent in parents: + if parent.type == "ClassDef": + full_name = f"{parent.name}.{full_name}" + break + return full_name def generate_candidates(source_code_path: Path) -> list[str]: diff --git a/codeflash/models/models.py b/codeflash/models/models.py index 244493233..27f36ca67 100644 --- a/codeflash/models/models.py +++ b/codeflash/models/models.py @@ -16,7 +16,11 @@ 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 extract_dependent_function, generate_candidates +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.test_results import TestResults, TestType @@ -322,18 +326,19 @@ def _fetch_function_coverages( 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=function_name, - coverage=coverage_data[function_name]["summary"]["percent_covered"], - executed_lines=coverage_data[function_name]["executed_lines"], - unexecuted_lines=coverage_data[function_name]["missing_lines"], - executed_branches=coverage_data[function_name]["executed_branches"], - unexecuted_branches=coverage_data[function_name]["missing_branches"], + 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=function_name, + name=resolved_name, coverage=0, executed_lines=[], unexecuted_lines=[], diff --git a/codeflash/optimization/function_optimizer.py b/codeflash/optimization/function_optimizer.py new file mode 100644 index 000000000..7b067a094 --- /dev/null +++ b/codeflash/optimization/function_optimizer.py @@ -0,0 +1,1152 @@ +from __future__ import annotations + +import ast +import concurrent.futures +import os +import shutil +import subprocess +import time +import uuid +from collections import defaultdict +from pathlib import Path +from typing import TYPE_CHECKING + +import isort +import libcst as cst +from rich.console import Group +from rich.panel import Panel +from rich.syntax import Syntax +from rich.tree import Tree + +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, + module_name_from_file_path, +) +from codeflash.code_utils.config_consts import ( + INDIVIDUAL_TESTCASE_TIMEOUT, + N_CANDIDATES, + N_TESTS_TO_GENERATE, + TOTAL_LOOPING_TIME, +) +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.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 +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 ( + BestOptimization, + CodeOptimizationContext, + FunctionCalledInTest, + FunctionParent, + GeneratedTests, + GeneratedTestsList, + OptimizationSet, + OptimizedCandidateResult, + OriginalCodeBaseline, + TestFile, + TestFiles, + TestingMode, +) +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 +from codeflash.telemetry.posthog_cf import ph +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_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 + +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 + + +class FunctionOptimizer: + def __init__( + self, + function_to_optimize: FunctionToOptimize, + test_cfg: TestConfig, + function_to_optimize_source_code: str = "", + function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, + function_to_optimize_ast: ast.FunctionDef | None = None, + aiservice_client: AiServiceClient | None = None, + args: Namespace | None = None, + ) -> None: + self.project_root = test_cfg.project_root_path + self.test_cfg = test_cfg + self.aiservice_client = aiservice_client if aiservice_client else AiServiceClient() + self.function_to_optimize = function_to_optimize + self.function_to_optimize_source_code = ( + function_to_optimize_source_code + if function_to_optimize_source_code + else function_to_optimize.file_path.read_text(encoding="utf8") + ) + if not function_to_optimize_ast: + original_module_ast = ast.parse(function_to_optimize_source_code) + self.function_to_optimize_ast = get_first_top_level_function_or_method_ast( + function_to_optimize.function_name, function_to_optimize.parents, original_module_ast + ) + else: + self.function_to_optimize_ast = function_to_optimize_ast + self.function_to_tests = function_to_tests if function_to_tests else {} + + 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) + + 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}") + ph("cli-optimize-function-start", {"function_trace_id": self.function_trace_id}) + self.cleanup_leftover_test_return_values() + file_name_from_test_module_name.cache_clear() + ctx_result = self.get_code_optimization_context() + if not is_successful(ctx_result): + return Failure(ctx_result.failure()) + code_context: CodeOptimizationContext = ctx_result.unwrap() + original_helper_code: dict[Path, str] = {} + helper_function_paths = {hf.file_path for hf in code_context.helper_functions} + for helper_function_path in helper_function_paths: + with helper_function_path.open(encoding="utf8") as f: + helper_code = f.read() + original_helper_code[helper_function_path] = helper_code + + 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" + ) + for test_index in range(N_TESTS_TO_GENERATE) + ] + generated_perf_test_paths = [ + get_test_file_path( + self.test_cfg.tests_root, self.function_to_optimize.function_name, test_index, test_type="perf" + ) + for test_index in range(N_TESTS_TO_GENERATE) + ] + + with progress_bar( + f"Generating new tests and optimizations for function {self.function_to_optimize.function_name}", + transient=True, + ): + generated_results = self.generate_tests_and_optimizations( + 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, + generated_test_paths=generated_test_paths, + generated_perf_test_paths=generated_perf_test_paths, + run_experiment=should_run_experiment, + ) + + if not is_successful(generated_results): + return Failure(generated_results.failure()) + generated_tests: GeneratedTestsList + optimizations_set: OptimizationSet + generated_tests, function_to_concolic_tests, concolic_test_str, optimizations_set = generated_results.unwrap() + count_tests = len(generated_tests.generated_tests) + if concolic_test_str: + count_tests += 1 + + for i, generated_test in enumerate(generated_tests.generated_tests): + with generated_test.behavior_file_path.open("w", encoding="utf8") as f: + f.write(generated_test.instrumented_behavior_test_source) + with generated_test.perf_file_path.open("w", encoding="utf8") as f: + f.write(generated_test.instrumented_perf_test_source) + self.test_files.add( + TestFile( + instrumented_behavior_file_path=generated_test.behavior_file_path, + benchmarking_file_path=generated_test.perf_file_path, + original_file_path=None, + original_source=generated_test.generated_original_test_source, + test_type=TestType.GENERATED_REGRESSION, + tests_in_file=None, # This is currently unused. We can discover the tests in the file if needed. + ) + ) + logger.info(f"Generated test {i + 1}/{count_tests}:") + code_print(generated_test.generated_original_test_source) + if concolic_test_str: + logger.info(f"Generated test {count_tests}/{count_tests}:") + code_print(concolic_test_str) + + function_to_optimize_qualified_name = self.function_to_optimize.qualified_name + function_to_all_tests = { + key: self.function_to_tests.get(key, []) + function_to_concolic_tests.get(key, []) + for key in set(self.function_to_tests) | set(function_to_concolic_tests) + } + instrumented_unittests_created_for_function = self.instrument_existing_tests(function_to_all_tests) + + # Get a dict of file_path_to_classes of fto and helpers_of_fto + file_path_to_helper_classes = defaultdict(set) + for function_source in code_context.helper_functions: + if ( + function_source.qualified_name != self.function_to_optimize.qualified_name + and "." in function_source.qualified_name + ): + file_path_to_helper_classes[function_source.file_path].add(function_source.qualified_name.split(".")[0]) + + baseline_result = self.establish_original_code_baseline( # this needs better typing + code_context=code_context, + original_helper_code=original_helper_code, + file_path_to_helper_classes=file_path_to_helper_classes, + ) + + console.rule() + paths_to_cleanup = ( + generated_test_paths + generated_perf_test_paths + list(instrumented_unittests_created_for_function) + ) + + if not is_successful(baseline_result): + cleanup_paths(paths_to_cleanup) + return Failure(baseline_result.failure()) + + original_code_baseline, test_functions_to_remove = baseline_result.unwrap() + if isinstance(original_code_baseline, OriginalCodeBaseline) and not coverage_critic( + original_code_baseline.coverage_results, self.args.test_framework + ): + cleanup_paths(paths_to_cleanup) + return Failure("The threshold for test coverage was not met.") + + best_optimization = None + + for u, candidates in enumerate([optimizations_set.control, optimizations_set.experiment]): + if candidates is None: + continue + + best_optimization = self.determine_best_candidate( + candidates=candidates, + code_context=code_context, + original_code_baseline=original_code_baseline, + original_helper_code=original_helper_code, + file_path_to_helper_classes=file_path_to_helper_classes, + ) + ph("cli-optimize-function-finished", {"function_trace_id": self.function_trace_id}) + + generated_tests = remove_functions_from_generated_tests( + generated_tests=generated_tests, test_functions_to_remove=test_functions_to_remove + ) + + if best_optimization: + logger.info("Best candidate:") + code_print(best_optimization.candidate.source_code) + console.print( + Panel( + best_optimization.candidate.explanation, title="Best Candidate Explanation", border_style="blue" + ) + ) + explanation = Explanation( + raw_explanation_message=best_optimization.candidate.explanation, + winning_behavioral_test_results=best_optimization.winning_behavioral_test_results, + winning_benchmarking_test_results=best_optimization.winning_benchmarking_test_results, + original_runtime_ns=original_code_baseline.runtime, + best_runtime_ns=best_optimization.runtime, + function_name=function_to_optimize_qualified_name, + file_path=self.function_to_optimize.file_path, + ) + + self.log_successful_optimization(explanation, generated_tests) + + self.replace_function_and_helpers_with_optimized_code( + code_context=code_context, optimized_code=best_optimization.candidate.source_code + ) + + new_code, new_helper_code = self.reformat_code_and_helpers( + code_context.helper_functions, explanation.file_path, self.function_to_optimize_source_code + ) + + existing_tests = existing_tests_source_for( + self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root), + function_to_all_tests, + tests_root=self.test_cfg.tests_root, + ) + + original_code_combined = original_helper_code.copy() + original_code_combined[explanation.file_path] = self.function_to_optimize_source_code + new_code_combined = new_helper_code.copy() + new_code_combined[explanation.file_path] = new_code + if not self.args.no_pr: + coverage_message = ( + original_code_baseline.coverage_results.build_message() + if original_code_baseline.coverage_results + else "Coverage data not available" + ) + generated_tests_str = "\n\n".join( + [test.generated_original_test_source for test in generated_tests.generated_tests] + ) + if concolic_test_str: + generated_tests_str += "\n\n" + concolic_test_str + + check_create_pr( + original_code=original_code_combined, + new_code=new_code_combined, + explanation=explanation, + existing_tests_source=existing_tests, + generated_original_test_source=generated_tests_str, + function_trace_id=self.function_trace_id, + coverage_message=coverage_message, + git_remote=self.args.git_remote, + ) + if self.args.all or env_utils.get_pr_number(): + self.write_code_and_helpers( + self.function_to_optimize_source_code, + original_helper_code, + self.function_to_optimize.file_path, + ) + for generated_test_path in generated_test_paths: + generated_test_path.unlink(missing_ok=True) + for generated_perf_test_path in generated_perf_test_paths: + generated_perf_test_path.unlink(missing_ok=True) + for test_paths in instrumented_unittests_created_for_function: + test_paths.unlink(missing_ok=True) + for fn in function_to_concolic_tests: + for test in function_to_concolic_tests[fn]: + if not test.tests_in_file.test_file.parent.exists(): + logger.warning( + f"Concolic test directory {test.tests_in_file.test_file.parent} does not exist so could not be deleted." + ) + shutil.rmtree(test.tests_in_file.test_file.parent, ignore_errors=True) + break # need to delete only one test directory + + if not best_optimization: + return Failure(f"No best optimizations found for function {self.function_to_optimize.qualified_name}") + return Success(best_optimization) + + def determine_best_candidate( + self, + *, + candidates: list[OptimizedCandidate], + code_context: CodeOptimizationContext, + original_code_baseline: OriginalCodeBaseline, + original_helper_code: dict[Path, str], + file_path_to_helper_classes: dict[Path, set[str]], + ) -> BestOptimization | None: + best_optimization: BestOptimization | None = None + best_runtime_until_now = original_code_baseline.runtime + + speedup_ratios: dict[str, float | None] = {} + optimized_runtimes: dict[str, float | None] = {} + is_correct = {} + + logger.info( + f"Determining best optimization candidate (out of {len(candidates)}) for " + 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." + ) + 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() + + 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 + + self.aiservice_client.log_results( + function_trace_id=self.function_trace_id, + speedup_ratio=speedup_ratios, + original_runtime=original_code_baseline.runtime, + optimized_runtime=optimized_runtimes, + is_correct=is_correct, + ) + return best_optimization + + def log_successful_optimization(self, explanation: Explanation, generated_tests: GeneratedTestsList) -> None: + explanation_panel = Panel( + f"⚡️ Optimization successful! 📄 {self.function_to_optimize.qualified_name} in {explanation.file_path}\n" + f"📈 {explanation.perf_improvement_line}\n" + f"Explanation: \n{explanation.to_console_string()}", + title="Optimization Summary", + border_style="green", + ) + + if self.args.no_pr: + tests_panel = Panel( + Syntax( + "\n".join([test.generated_original_test_source for test in generated_tests.generated_tests]), + "python", + line_numbers=True, + ), + title="Validated Tests", + border_style="blue", + ) + + console.print(Group(explanation_panel, tests_panel)) + console.print(explanation_panel) + + ph( + "cli-optimize-success", + { + "function_trace_id": self.function_trace_id, + "speedup_x": explanation.speedup_x, + "speedup_pct": explanation.speedup_pct, + "best_runtime": explanation.best_runtime_ns, + "original_runtime": explanation.original_runtime_ns, + "winning_test_results": { + tt.to_name(): v + for tt, v in explanation.winning_behavioral_test_results.get_test_pass_fail_report_by_type().items() + }, + }, + ) + + @staticmethod + def write_code_and_helpers(original_code: str, original_helper_code: dict[Path, str], path: Path) -> None: + with path.open("w", encoding="utf8") as f: + f.write(original_code) + for module_abspath in original_helper_code: + with Path(module_abspath).open("w", encoding="utf8") as f: + f.write(original_helper_code[module_abspath]) + + def reformat_code_and_helpers( + self, helper_functions: list[FunctionSource], path: Path, original_code: str + ) -> tuple[str, dict[Path, str]]: + should_sort_imports = not self.args.disable_imports_sorting + if should_sort_imports and isort.code(original_code) != original_code: + should_sort_imports = False + + new_code = format_code(self.args.formatter_cmds, path) + if should_sort_imports: + new_code = sort_imports(new_code) + + new_helper_code: dict[Path, str] = {} + helper_functions_paths = {hf.file_path for hf in helper_functions} + for module_abspath in helper_functions_paths: + formatted_helper_code = format_code(self.args.formatter_cmds, module_abspath) + if should_sort_imports: + formatted_helper_code = sort_imports(formatted_helper_code) + new_helper_code[module_abspath] = formatted_helper_code + + return new_code, new_helper_code + + def replace_function_and_helpers_with_optimized_code( + self, code_context: CodeOptimizationContext, optimized_code: str + ) -> bool: + did_update = False + read_writable_functions_by_file_path = defaultdict(set) + read_writable_functions_by_file_path[self.function_to_optimize.file_path].add( + self.function_to_optimize.qualified_name + ) + for helper_function in code_context.helper_functions: + if helper_function.jedi_definition.type != "class": + read_writable_functions_by_file_path[helper_function.file_path].add(helper_function.qualified_name) + for module_abspath, qualified_names in read_writable_functions_by_file_path.items(): + did_update |= replace_function_definitions_in_module( + function_names=list(qualified_names), + optimized_code=optimized_code, + module_abspath=module_abspath, + preexisting_objects=code_context.preexisting_objects, + project_root_path=self.project_root, + ) + 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 + ) + except ValueError as e: + return Failure(str(e)) + + return Success( + CodeOptimizationContext( + 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 + preexisting_objects=new_code_ctx.preexisting_objects, + ) + ) + + @staticmethod + def cleanup_leftover_test_return_values() -> None: + # remove leftovers from previous run + get_run_tmp_file(Path("test_return_values_0.bin")).unlink(missing_ok=True) + get_run_tmp_file(Path("test_return_values_0.sqlite")).unlink(missing_ok=True) + + def instrument_existing_tests(self, function_to_all_tests: dict[str, list[FunctionCalledInTest]]) -> set[Path]: + existing_test_files_count = 0 + replay_test_files_count = 0 + concolic_coverage_test_files_count = 0 + unique_instrumented_test_files = set() + + func_qualname = self.function_to_optimize.qualified_name_with_modules_from_root(self.project_root) + if func_qualname not in function_to_all_tests: + logger.info(f"Did not find any pre-existing tests for '{func_qualname}', will only use generated tests.") + console.rule() + else: + test_file_invocation_positions = defaultdict(list[FunctionCalledInTest]) + for tests_in_file in function_to_all_tests.get(func_qualname): + test_file_invocation_positions[ + (tests_in_file.tests_in_file.test_file, tests_in_file.tests_in_file.test_type) + ].append(tests_in_file) + for (test_file, test_type), tests_in_file_list in test_file_invocation_positions.items(): + path_obj_test_file = Path(test_file) + if test_type == TestType.EXISTING_UNIT_TEST: + existing_test_files_count += 1 + elif test_type == TestType.REPLAY_TEST: + replay_test_files_count += 1 + elif test_type == TestType.CONCOLIC_COVERAGE_TEST: + concolic_coverage_test_files_count += 1 + else: + msg = f"Unexpected test type: {test_type}" + raise ValueError(msg) + success, injected_behavior_test = inject_profiling_into_existing_test( + mode=TestingMode.BEHAVIOR, + test_path=path_obj_test_file, + call_positions=[test.position for test in tests_in_file_list], + function_to_optimize=self.function_to_optimize, + tests_project_root=self.test_cfg.tests_project_rootdir, + test_framework=self.args.test_framework, + ) + if not success: + continue + success, injected_perf_test = inject_profiling_into_existing_test( + mode=TestingMode.PERFORMANCE, + test_path=path_obj_test_file, + call_positions=[test.position for test in tests_in_file_list], + function_to_optimize=self.function_to_optimize, + tests_project_root=self.test_cfg.tests_project_rootdir, + test_framework=self.args.test_framework, + ) + if not success: + continue + # TODO: this naming logic should be moved to a function and made more standard + new_behavioral_test_path = Path( + f"{os.path.splitext(test_file)[0]}__perfinstrumented{os.path.splitext(test_file)[1]}" + ) + new_perf_test_path = Path( + f"{os.path.splitext(test_file)[0]}__perfonlyinstrumented{os.path.splitext(test_file)[1]}" + ) + if injected_behavior_test is not None: + with new_behavioral_test_path.open("w", encoding="utf8") as _f: + _f.write(injected_behavior_test) + else: + msg = "injected_behavior_test is None" + raise ValueError(msg) + if injected_perf_test is not None: + with new_perf_test_path.open("w", encoding="utf8") as _f: + _f.write(injected_perf_test) + + 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( + instrumented_behavior_file_path=new_behavioral_test_path, + benchmarking_file_path=new_perf_test_path, + original_source=None, + original_file_path=Path(test_file), + test_type=test_type, + 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" + f"{'s' if replay_test_files_count != 1 else ''}, and " + f"{concolic_coverage_test_files_count} concolic coverage test file" + f"{'s' if concolic_coverage_test_files_count != 1 else ''} for {func_qualname}" + ) + return unique_instrumented_test_files + + def generate_tests_and_optimizations( + self, + code_to_optimize_with_helpers: str, + read_writable_code: str, + read_only_context_code: str, + helper_functions: list[FunctionSource], + generated_test_paths: list[Path], + generated_perf_test_paths: list[Path], + run_experiment: bool = False, + ) -> Result[tuple[GeneratedTestsList, dict[str, list[FunctionCalledInTest]], OptimizationSet], str]: + assert len(generated_test_paths) == N_TESTS_TO_GENERATE + max_workers = N_TESTS_TO_GENERATE + 2 if not run_experiment else N_TESTS_TO_GENERATE + 3 + console.rule() + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + # Submit the test generation task as future + future_tests = self.generate_and_instrument_tests( + executor, + code_to_optimize_with_helpers, + [definition.fully_qualified_name for definition in helper_functions], + generated_test_paths, + generated_perf_test_paths, + ) + future_optimization_candidates = executor.submit( + self.aiservice_client.optimize_python_code, + read_writable_code, + read_only_context_code, + self.function_trace_id[:-4] + "EXP0" if run_experiment else self.function_trace_id, + N_CANDIDATES, + ExperimentMetadata(id=self.experiment_id, group="control") if run_experiment else None, + ) + future_candidates_exp = None + + future_concolic_tests = executor.submit( + generate_concolic_tests, + self.test_cfg, + self.args, + self.function_to_optimize, + self.function_to_optimize_ast, + ) + futures = [*future_tests, future_optimization_candidates, future_concolic_tests] + if run_experiment: + future_candidates_exp = executor.submit( + self.local_aiservice_client.optimize_python_code, + read_writable_code, + read_only_context_code, + self.function_trace_id[:-4] + "EXP1", + N_CANDIDATES, + ExperimentMetadata(id=self.experiment_id, group="experiment"), + ) + futures.append(future_candidates_exp) + + # Wait for all futures to complete + concurrent.futures.wait(futures) + + # Retrieve results + candidates: list[OptimizedCandidate] = future_optimization_candidates.result() + if not candidates: + return Failure(f"/!\\ NO OPTIMIZATIONS GENERATED for {self.function_to_optimize.function_name}") + + candidates_experiment = future_candidates_exp.result() if future_candidates_exp else None + + # Process test generation results + + tests: list[GeneratedTests] = [] + for future in future_tests: + res = future.result() + if res: + ( + generated_test_source, + instrumented_behavior_test_source, + instrumented_perf_test_source, + test_behavior_path, + test_perf_path, + ) = res + tests.append( + GeneratedTests( + generated_original_test_source=generated_test_source, + instrumented_behavior_test_source=instrumented_behavior_test_source, + instrumented_perf_test_source=instrumented_perf_test_source, + behavior_file_path=test_behavior_path, + perf_file_path=test_perf_path, + ) + ) + if not tests: + logger.warning(f"Failed to generate and instrument tests for {self.function_to_optimize.function_name}") + return Failure(f"/!\\ NO TESTS GENERATED for {self.function_to_optimize.function_name}") + function_to_concolic_tests, concolic_test_str = future_concolic_tests.result() + logger.info(f"Generated {len(tests)} tests for {self.function_to_optimize.function_name}") + console.rule() + generated_tests = GeneratedTestsList(generated_tests=tests) + + return Success( + ( + generated_tests, + function_to_concolic_tests, + concolic_test_str, + OptimizationSet(control=candidates, experiment=candidates_experiment), + ) + ) + + def establish_original_code_baseline( + self, + code_context: CodeOptimizationContext, + original_helper_code: dict[Path, str], + file_path_to_helper_classes: dict[Path, set[str]], + ) -> Result[tuple[OriginalCodeBaseline, list[str]], str]: + # 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"] + success = True + + test_env = os.environ.copy() + test_env["CODEFLASH_TEST_ITERATION"] = "0" + test_env["CODEFLASH_TRACER_DISABLE"] = "1" + test_env["CODEFLASH_LOOP_INDEX"] = "0" + if "PYTHONPATH" not in test_env: + test_env["PYTHONPATH"] = str(self.args.project_root) + else: + test_env["PYTHONPATH"] += os.pathsep + str(self.args.project_root) + + coverage_results = None + # Instrument codeflash capture + try: + instrument_codeflash_capture( + self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root + ) + behavioral_results, coverage_results = self.run_and_parse_tests( + testing_type=TestingMode.BEHAVIOR, + test_env=test_env, + test_files=self.test_files, + optimization_iteration=0, + testing_time=TOTAL_LOOPING_TIME, + enable_coverage=test_framework == "pytest", + code_context=code_context, + ) + 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 not behavioral_results: + logger.warning( + f"Couldn't run any tests for original function {self.function_to_optimize.function_name}. SKIPPING OPTIMIZING THIS FUNCTION." + ) + 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 + ): + return Failure("The threshold for test coverage was not met.") + if test_framework == "pytest": + benchmarking_results, _ = self.run_and_parse_tests( + testing_type=TestingMode.PERFORMANCE, + test_env=test_env, + test_files=self.test_files, + optimization_iteration=0, + testing_time=TOTAL_LOOPING_TIME, + enable_coverage=False, + code_context=code_context, + ) + + else: + benchmarking_results = TestResults() + start_time: float = time.time() + for i in range(100): + if i >= 5 and time.time() - start_time >= TOTAL_LOOPING_TIME * 1.5: + # * 1.5 to give unittest a bit more time to run + break + test_env["CODEFLASH_LOOP_INDEX"] = str(i + 1) + unittest_loop_results, _ = self.run_and_parse_tests( + testing_type=TestingMode.PERFORMANCE, + test_env=test_env, + test_files=self.test_files, + optimization_iteration=0, + testing_time=TOTAL_LOOPING_TIME, + enable_coverage=False, + code_context=code_context, + unittest_loop_index=i + 1, + ) + benchmarking_results.merge(unittest_loop_results) + + console.print( + TestResults.report_to_tree( + behavioral_results.get_test_pass_fail_report_by_type(), + title="Overall test results for original code", + ) + ) + console.rule() + + + total_timing = benchmarking_results.total_passed_runtime() # caution: doesn't handle the loop index + functions_to_remove = [ + result.id.test_function_name + for result in behavioral_results + if (result.test_type == TestType.GENERATED_REGRESSION and not result.did_pass) + ] + if total_timing == 0: + logger.warning( + "The overall summed benchmark runtime of the original function is 0, couldn't run tests." + ) + console.rule() + success = False + if not total_timing: + logger.warning("Failed to run the tests for the original function, skipping optimization") + console.rule() + success = False + if not success: + return Failure("Failed to establish a baseline for the original code.") + + loop_count = max([int(result.loop_index) for result in benchmarking_results.test_results]) + logger.info( + f"Original code summed runtime measured over {loop_count} loop{'s' if loop_count > 1 else ''}: " + f"{humanize_runtime(total_timing)} per full loop" + ) + console.rule() + logger.debug(f"Total original code runtime (ns): {total_timing}") + return Success( + ( + OriginalCodeBaseline( + behavioral_test_results=behavioral_results, + benchmarking_test_results=benchmarking_results, + runtime=total_timing, + coverage_results=coverage_results, + ), + functions_to_remove, + ) + ) + + def run_optimized_candidate( + self, + *, + optimization_candidate_index: int, + baseline_results: OriginalCodeBaseline, + original_helper_code: dict[Path, str], + 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" + test_env["CODEFLASH_TEST_ITERATION"] = str(optimization_candidate_index) + test_env["CODEFLASH_TRACER_DISABLE"] = "1" + if "PYTHONPATH" not in test_env: + test_env["PYTHONPATH"] = str(self.project_root) + else: + 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 = {} + for module_abspath in original_helper_code: + candidate_helper_code[module_abspath] = Path(module_abspath).read_text("utf-8") + try: + instrument_codeflash_capture( + self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root + ) + + candidate_behavior_results, _ = self.run_and_parse_tests( + testing_type=TestingMode.BEHAVIOR, + test_env=test_env, + test_files=self.test_files, + optimization_iteration=optimization_candidate_index, + testing_time=TOTAL_LOOPING_TIME, + enable_coverage=False, + ) + # Remove instrumentation + finally: + self.write_code_and_helpers( + candidate_fto_code, candidate_helper_code, self.function_to_optimize.file_path + ) + console.print( + TestResults.report_to_tree( + candidate_behavior_results.get_test_pass_fail_report_by_type(), + title="Behavioral Test Results for candidate", + ) + ) + console.rule() + + if compare_test_results(baseline_results.behavioral_test_results, candidate_behavior_results): + logger.info("Test results matched!") + console.rule() + else: + logger.info("Test results did not match the test results of the original code.") + console.rule() + return Failure("Test results did not match the test results of the original code.") + + if test_framework == "pytest": + candidate_benchmarking_results, _ = self.run_and_parse_tests( + testing_type=TestingMode.PERFORMANCE, + test_env=test_env, + test_files=self.test_files, + optimization_iteration=optimization_candidate_index, + testing_time=TOTAL_LOOPING_TIME, + enable_coverage=False, + ) + loop_count = ( + max(all_loop_indices) + if ( + all_loop_indices := { + result.loop_index for result in candidate_benchmarking_results.test_results + } + ) + else 0 + ) + + else: + candidate_benchmarking_results = TestResults() + start_time: float = time.time() + loop_count = 0 + for i in range(100): + if i >= 5 and time.time() - start_time >= TOTAL_LOOPING_TIME * 1.5: + # * 1.5 to give unittest a bit more time to run + break + test_env["CODEFLASH_LOOP_INDEX"] = str(i + 1) + unittest_loop_results, cov = self.run_and_parse_tests( + testing_type=TestingMode.PERFORMANCE, + test_env=test_env, + test_files=self.test_files, + optimization_iteration=optimization_candidate_index, + testing_time=TOTAL_LOOPING_TIME, + unittest_loop_index=i + 1, + ) + loop_count = i + 1 + candidate_benchmarking_results.merge(unittest_loop_results) + + if (total_candidate_timing := candidate_benchmarking_results.total_passed_runtime()) == 0: + logger.warning("The overall test runtime of the optimized function is 0, couldn't run tests.") + console.rule() + + logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") + return Success( + OptimizedCandidateResult( + max_loop_count=loop_count, + best_test_runtime=total_candidate_timing, + behavior_test_results=candidate_behavior_results, + benchmarking_test_results=candidate_benchmarking_results, + optimization_candidate_index=optimization_candidate_index, + total_candidate_timing=total_candidate_timing, + ) + ) + + def run_and_parse_tests( + self, + testing_type: TestingMode, + test_env: dict[str, str], + test_files: TestFiles, + optimization_iteration: int, + testing_time: float = TOTAL_LOOPING_TIME, + *, + enable_coverage: bool = False, + pytest_min_loops: int = 5, + pytest_max_loops: int = 100_000, + code_context: CodeOptimizationContext | None = None, + unittest_loop_index: int | None = None, + ) -> tuple[TestResults, CoverageData | None]: + coverage_database_file = None + try: + if testing_type == TestingMode.BEHAVIOR: + 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.PERFORMANCE: + result_file_path, run_result = run_benchmarking_tests( + test_files, + cwd=self.project_root, + test_env=test_env, + 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: + 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' + ) + 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"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, + ) + return results, coverage_results + + def generate_and_instrument_tests( + self, + executor: concurrent.futures.ThreadPoolExecutor, + source_code_being_tested: str, + helper_function_names: list[str], + generated_test_paths: list[Path], + generated_perf_test_paths: list[Path], + ) -> list[concurrent.futures.Future]: + return [ + executor.submit( + generate_tests, + self.aiservice_client, + source_code_being_tested, + self.function_to_optimize, + helper_function_names, + Path(self.original_module_path), + self.test_cfg, + INDIVIDUAL_TESTCASE_TIMEOUT, + self.function_trace_id, + test_index, + test_path, + test_perf_path, + ) + for test_index, (test_path, test_perf_path) in enumerate( + zip(generated_test_paths, generated_perf_test_paths) + ) + ] + diff --git a/codeflash/optimization/optimizer.py b/codeflash/optimization/optimizer.py index a50436bd6..cae78a153 100644 --- a/codeflash/optimization/optimizer.py +++ b/codeflash/optimization/optimizer.py @@ -1,85 +1,32 @@ from __future__ import annotations import ast -import concurrent.futures import os import shutil -import subprocess import tempfile -import time -import uuid -from collections import defaultdict from pathlib import Path from typing import TYPE_CHECKING -import isort -import libcst as cst -from rich.console import Group -from rich.panel import Panel -from rich.syntax import Syntax -from rich.tree import Tree - from codeflash.api.aiservice import AiServiceClient, LocalAiServiceClient -from codeflash.cli_cmds.console import code_print, console, logger, progress_bar +from codeflash.cli_cmds.console import console, logger 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 normalize_code, normalize_node, replace_function_definitions_in_module -from codeflash.code_utils.code_utils import ( - cleanup_paths, - file_name_from_test_module_name, - get_run_tmp_file, - module_name_from_file_path, -) -from codeflash.code_utils.config_consts import ( - INDIVIDUAL_TESTCASE_TIMEOUT, - N_CANDIDATES, - N_TESTS_TO_GENERATE, - TOTAL_LOOPING_TIME, -) -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.remove_generated_tests import remove_functions_from_generated_tests +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 -from codeflash.code_utils.time_utils import humanize_runtime -from codeflash.context import code_context_extractor from codeflash.discovery.discover_unit_tests import discover_unit_tests -from codeflash.discovery.functions_to_optimize import FunctionToOptimize, get_functions_to_optimize -from codeflash.either import Failure, Success, is_successful -from codeflash.models.ExperimentMetadata import ExperimentMetadata -from codeflash.models.models import ( - BestOptimization, - CodeOptimizationContext, - FunctionCalledInTest, - FunctionParent, - GeneratedTests, - GeneratedTestsList, - OptimizationSet, - OptimizedCandidateResult, - OriginalCodeBaseline, - TestFile, - TestFiles, - TestingMode, - ValidCode, -) -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 +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.optimization.function_optimizer import FunctionOptimizer from codeflash.telemetry.posthog_cf import ph -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_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 TestConfig, get_test_file_path -from codeflash.verification.verifier import generate_tests +from codeflash.verification.test_results import TestType +from codeflash.verification.verification_utils import TestConfig if TYPE_CHECKING: from argparse import Namespace - from codeflash.either import Result - from codeflash.models.models import CoverageData, FunctionSource, OptimizedCandidate + from codeflash.discovery.functions_to_optimize import FunctionToOptimize + from codeflash.models.models import FunctionCalledInTest class Optimizer: @@ -98,7 +45,23 @@ def __init__(self, args: Namespace) -> None: 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=[]) + def create_function_optimizer( + self, + function_to_optimize: FunctionToOptimize, + function_to_optimize_ast: ast.FunctionDef | None = None, + function_to_tests: dict[str, list[FunctionCalledInTest]] | None = None, + function_to_optimize_source_code: str | None = "", + ) -> FunctionOptimizer: + return FunctionOptimizer( + function_to_optimize=function_to_optimize, + test_cfg=self.test_cfg, + function_to_optimize_source_code=function_to_optimize_source_code, + function_to_tests=function_to_tests, + function_to_optimize_ast=function_to_optimize_ast, + aiservice_client=self.aiservice_client, + args=self.args, + + ) def run(self) -> None: ph("cli-optimize-run-start") @@ -106,7 +69,7 @@ def run(self) -> None: console.rule() if not env_utils.ensure_codeflash_api_key(): return - + function_optimizer = None file_to_funcs_to_optimize: dict[Path, list[FunctionToOptimize]] num_optimizable_functions: int (file_to_funcs_to_optimize, num_optimizable_functions) = get_functions_to_optimize( @@ -198,10 +161,10 @@ def run(self) -> None: ) continue - best_optimization = self.optimize_function( - function_to_optimize, function_to_optimize_ast, function_to_tests, validated_original_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 ) - self.test_files = TestFiles(test_files=[]) + best_optimization = function_optimizer.optimize_function() if is_successful(best_optimization): optimizations_found += 1 else: @@ -214,1100 +177,20 @@ def run(self) -> None: elif self.args.all: logger.info("✨ All functions have been optimized! ✨") finally: - for test_file in self.test_files.get_by_type(TestType.GENERATED_REGRESSION).test_files: - test_file.instrumented_behavior_file_path.unlink(missing_ok=True) - test_file.benchmarking_file_path.unlink(missing_ok=True) - for test_file in self.test_files.get_by_type(TestType.EXISTING_UNIT_TEST).test_files: - test_file.instrumented_behavior_file_path.unlink(missing_ok=True) - test_file.benchmarking_file_path.unlink(missing_ok=True) - for test_file in self.test_files.get_by_type(TestType.CONCOLIC_COVERAGE_TEST).test_files: - test_file.instrumented_behavior_file_path.unlink(missing_ok=True) + if function_optimizer: + for test_file in function_optimizer.test_files.get_by_type(TestType.GENERATED_REGRESSION).test_files: + test_file.instrumented_behavior_file_path.unlink(missing_ok=True) + test_file.benchmarking_file_path.unlink(missing_ok=True) + for test_file in function_optimizer.test_files.get_by_type(TestType.EXISTING_UNIT_TEST).test_files: + test_file.instrumented_behavior_file_path.unlink(missing_ok=True) + test_file.benchmarking_file_path.unlink(missing_ok=True) + for test_file in function_optimizer.test_files.get_by_type(TestType.CONCOLIC_COVERAGE_TEST).test_files: + 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 hasattr(get_run_tmp_file, "tmpdir"): get_run_tmp_file.tmpdir.cleanup() - if self.test_cfg.concolic_test_root_dir: - shutil.rmtree(self.test_cfg.concolic_test_root_dir, ignore_errors=True) - - def optimize_function( - self, - function_to_optimize: FunctionToOptimize, - function_to_optimize_ast: ast.FunctionDef, - function_to_tests: dict[str, list[FunctionCalledInTest]], - validated_original_code: dict[Path, ValidCode], - ) -> Result[BestOptimization, str]: - should_run_experiment = self.experiment_id is not None - function_trace_id: str = str(uuid.uuid4()) - logger.debug(f"Function Trace ID: {function_trace_id}") - ph("cli-optimize-function-start", {"function_trace_id": function_trace_id}) - self.cleanup_leftover_test_return_values() - file_name_from_test_module_name.cache_clear() - ctx_result = self.get_code_optimization_context( - function_to_optimize, - self.args.project_root, - validated_original_code[function_to_optimize.file_path].source_code, - ) - if not is_successful(ctx_result): - return Failure(ctx_result.failure()) - code_context: CodeOptimizationContext = ctx_result.unwrap() - original_helper_code: dict[Path, str] = {} - helper_function_paths = {hf.file_path for hf in code_context.helper_functions} - for helper_function_path in helper_function_paths: - with helper_function_path.open(encoding="utf8") as f: - helper_code = f.read() - original_helper_code[helper_function_path] = helper_code - - logger.info("Code to be optimized:") - code_print(code_context.read_writable_code) - - original_module_path = module_name_from_file_path(function_to_optimize.file_path, self.args.project_root) - - 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, - function_to_optimize.file_path, - self.args.project_root, - ) - - generated_test_paths = [ - get_test_file_path( - self.test_cfg.tests_root, function_to_optimize.function_name, test_index, test_type="unit" - ) - for test_index in range(N_TESTS_TO_GENERATE) - ] - generated_perf_test_paths = [ - get_test_file_path( - self.test_cfg.tests_root, function_to_optimize.function_name, test_index, test_type="perf" - ) - for test_index in range(N_TESTS_TO_GENERATE) - ] - - with progress_bar( - f"Generating new tests and optimizations for function {function_to_optimize.function_name}", transient=True - ): - generated_results = self.generate_tests_and_optimizations( - 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, - function_to_optimize=function_to_optimize, - helper_functions=code_context.helper_functions, - module_path=Path(original_module_path), - function_trace_id=function_trace_id, - generated_test_paths=generated_test_paths, - generated_perf_test_paths=generated_perf_test_paths, - function_to_optimize_ast=function_to_optimize_ast, - run_experiment=should_run_experiment, - ) - - if not is_successful(generated_results): - return Failure(generated_results.failure()) - generated_tests: GeneratedTestsList - optimizations_set: OptimizationSet - generated_tests, function_to_concolic_tests, concolic_test_str, optimizations_set = generated_results.unwrap() - count_tests = len(generated_tests.generated_tests) - if concolic_test_str: - count_tests += 1 - - for i, generated_test in enumerate(generated_tests.generated_tests): - with generated_test.behavior_file_path.open("w", encoding="utf8") as f: - f.write(generated_test.instrumented_behavior_test_source) - with generated_test.perf_file_path.open("w", encoding="utf8") as f: - f.write(generated_test.instrumented_perf_test_source) - self.test_files.add( - TestFile( - instrumented_behavior_file_path=generated_test.behavior_file_path, - benchmarking_file_path=generated_test.perf_file_path, - original_file_path=None, - original_source=generated_test.generated_original_test_source, - test_type=TestType.GENERATED_REGRESSION, - tests_in_file=None, # This is currently unused. We can discover the tests in the file if needed. - ) - ) - logger.info(f"Generated test {i + 1}/{count_tests}:") - code_print(generated_test.generated_original_test_source) - if concolic_test_str: - logger.info(f"Generated test {count_tests}/{count_tests}:") - code_print(concolic_test_str) - - function_to_optimize_qualified_name = function_to_optimize.qualified_name - function_to_all_tests = { - key: function_to_tests.get(key, []) + function_to_concolic_tests.get(key, []) - for key in set(function_to_tests) | set(function_to_concolic_tests) - } - instrumented_unittests_created_for_function = self.instrument_existing_tests( - function_to_optimize=function_to_optimize, function_to_tests=function_to_all_tests - ) - - # Get a dict of file_path_to_classes of fto and helpers_of_fto - file_path_to_helper_classes = defaultdict(set) - for function_source in code_context.helper_functions: - if ( - function_source.qualified_name != function_to_optimize.qualified_name - and "." in function_source.qualified_name - ): - file_path_to_helper_classes[function_source.file_path].add(function_source.qualified_name.split(".")[0]) - - baseline_result = self.establish_original_code_baseline( # this needs better typing - function_name=function_to_optimize_qualified_name, - function_file_path=function_to_optimize.file_path, - code_context=code_context, - function_to_optimize=function_to_optimize, - original_helper_code=original_helper_code, - file_path_to_helper_classes=file_path_to_helper_classes, - ) - - console.rule() - paths_to_cleanup = ( - generated_test_paths + generated_perf_test_paths + list(instrumented_unittests_created_for_function) - ) - - if not is_successful(baseline_result): - cleanup_paths(paths_to_cleanup) - return Failure(baseline_result.failure()) - - original_code_baseline, test_functions_to_remove = baseline_result.unwrap() - if isinstance(original_code_baseline, OriginalCodeBaseline) and not coverage_critic( - original_code_baseline.coverage_results, self.args.test_framework - ): - cleanup_paths(paths_to_cleanup) - return Failure("The threshold for test coverage was not met.") - - best_optimization = None - - for u, candidates in enumerate([optimizations_set.control, optimizations_set.experiment]): - if candidates is None: - continue - - best_optimization = self.determine_best_candidate( - candidates=candidates, - code_context=code_context, - function_to_optimize=function_to_optimize, - original_code=validated_original_code[function_to_optimize.file_path].source_code, - original_code_baseline=original_code_baseline, - original_helper_code=original_helper_code, - function_trace_id=function_trace_id[:-4] + f"EXP{u}" if should_run_experiment else function_trace_id, - file_path_to_helper_classes=file_path_to_helper_classes, - ) - ph("cli-optimize-function-finished", {"function_trace_id": function_trace_id}) - - generated_tests = remove_functions_from_generated_tests( - generated_tests=generated_tests, test_functions_to_remove=test_functions_to_remove - ) - - if best_optimization: - logger.info("Best candidate:") - code_print(best_optimization.candidate.source_code) - console.print( - Panel( - best_optimization.candidate.explanation, title="Best Candidate Explanation", border_style="blue" - ) - ) - explanation = Explanation( - raw_explanation_message=best_optimization.candidate.explanation, - winning_behavioral_test_results=best_optimization.winning_behavioral_test_results, - winning_benchmarking_test_results=best_optimization.winning_benchmarking_test_results, - original_runtime_ns=original_code_baseline.runtime, - best_runtime_ns=best_optimization.runtime, - function_name=function_to_optimize_qualified_name, - file_path=function_to_optimize.file_path, - ) - - self.log_successful_optimization(explanation, function_to_optimize, function_trace_id, generated_tests) - - self.replace_function_and_helpers_with_optimized_code( - code_context=code_context, - function_to_optimize_file_path=explanation.file_path, - optimized_code=best_optimization.candidate.source_code, - qualified_function_name=function_to_optimize_qualified_name, - ) - - new_code, new_helper_code = self.reformat_code_and_helpers( - code_context.helper_functions, - explanation.file_path, - validated_original_code[function_to_optimize.file_path].source_code, - ) - - existing_tests = existing_tests_source_for( - function_to_optimize.qualified_name_with_modules_from_root(self.args.project_root), - function_to_all_tests, - tests_root=self.test_cfg.tests_root, - ) - - original_code_combined = original_helper_code.copy() - original_code_combined[explanation.file_path] = validated_original_code[ - function_to_optimize.file_path - ].source_code - new_code_combined = new_helper_code.copy() - new_code_combined[explanation.file_path] = new_code - if not self.args.no_pr: - coverage_message = ( - original_code_baseline.coverage_results.build_message() - if original_code_baseline.coverage_results - else "Coverage data not available" - ) - generated_tests_str = "\n\n".join( - [test.generated_original_test_source for test in generated_tests.generated_tests] - ) - if concolic_test_str: - generated_tests_str += "\n\n" + concolic_test_str - - check_create_pr( - original_code=original_code_combined, - new_code=new_code_combined, - explanation=explanation, - existing_tests_source=existing_tests, - generated_original_test_source=generated_tests_str, - function_trace_id=function_trace_id, - coverage_message=coverage_message, - git_remote=self.args.git_remote, - ) - if self.args.all or env_utils.get_pr_number(): - self.write_code_and_helpers( - validated_original_code[function_to_optimize.file_path].source_code, - original_helper_code, - function_to_optimize.file_path, - ) - for generated_test_path in generated_test_paths: - generated_test_path.unlink(missing_ok=True) - for generated_perf_test_path in generated_perf_test_paths: - generated_perf_test_path.unlink(missing_ok=True) - for test_paths in instrumented_unittests_created_for_function: - test_paths.unlink(missing_ok=True) - for fn in function_to_concolic_tests: - for test in function_to_concolic_tests[fn]: - if not test.tests_in_file.test_file.parent.exists(): - logger.warning( - f"Concolic test directory {test.tests_in_file.test_file.parent} does not exist so could not be deleted." - ) - shutil.rmtree(test.tests_in_file.test_file.parent, ignore_errors=True) - break # need to delete only one test directory - - if not best_optimization: - return Failure(f"No best optimizations found for function {function_to_optimize.qualified_name}") - return Success(best_optimization) - - def determine_best_candidate( - self, - *, - candidates: list[OptimizedCandidate], - code_context: CodeOptimizationContext, - function_to_optimize: FunctionToOptimize, - original_code: str, - original_code_baseline: OriginalCodeBaseline, - original_helper_code: dict[Path, str], - function_trace_id: str, - file_path_to_helper_classes: dict[Path, set[str]], - ) -> BestOptimization | None: - best_optimization: BestOptimization | None = None - best_runtime_until_now = original_code_baseline.runtime - - speedup_ratios: dict[str, float | None] = {} - optimized_runtimes: dict[str, float | None] = {} - is_correct = {} - - logger.info( - f"Determining best optimization candidate (out of {len(candidates)}) for " - f"{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, - function_to_optimize_file_path=function_to_optimize.file_path, - optimized_code=candidate.source_code, - qualified_function_name=function_to_optimize.qualified_name, - ) - 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(original_code, original_helper_code, 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, - function_to_optimize=function_to_optimize, - 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:.3f}") - - 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:.3f}") - console.print(tree) - console.rule() - - self.write_code_and_helpers(original_code, original_helper_code, function_to_optimize.file_path) - except KeyboardInterrupt as e: - self.write_code_and_helpers(original_code, original_helper_code, function_to_optimize.file_path) - logger.exception(f"Optimization interrupted: {e}") - raise - - self.aiservice_client.log_results( - function_trace_id=function_trace_id, - speedup_ratio=speedup_ratios, - original_runtime=original_code_baseline.runtime, - optimized_runtime=optimized_runtimes, - is_correct=is_correct, - ) - return best_optimization - - def log_successful_optimization( - self, - explanation: Explanation, - function_to_optimize: FunctionToOptimize, - function_trace_id: str, - generated_tests: GeneratedTestsList, - ) -> None: - explanation_panel = Panel( - f"⚡️ Optimization successful! 📄 {function_to_optimize.qualified_name} in {explanation.file_path}\n" - f"📈 {explanation.perf_improvement_line}\n" - f"Explanation: \n{explanation.to_console_string()}", - title="Optimization Summary", - border_style="green", - ) - - if self.args.no_pr: - tests_panel = Panel( - Syntax( - "\n".join([test.generated_original_test_source for test in generated_tests.generated_tests]), - "python", - line_numbers=True, - ), - title="Validated Tests", - border_style="blue", - ) - - console.print(Group(explanation_panel, tests_panel)) - console.print(explanation_panel) - - ph( - "cli-optimize-success", - { - "function_trace_id": function_trace_id, - "speedup_x": explanation.speedup_x, - "speedup_pct": explanation.speedup_pct, - "best_runtime": explanation.best_runtime_ns, - "original_runtime": explanation.original_runtime_ns, - "winning_test_results": { - tt.to_name(): v - for tt, v in explanation.winning_behavioral_test_results.get_test_pass_fail_report_by_type().items() - }, - }, - ) - - @staticmethod - def write_code_and_helpers(original_code: str, original_helper_code: dict[Path, str], path: Path) -> None: - with path.open("w", encoding="utf8") as f: - f.write(original_code) - for module_abspath in original_helper_code: - with Path(module_abspath).open("w", encoding="utf8") as f: - f.write(original_helper_code[module_abspath]) - - def reformat_code_and_helpers( - self, helper_functions: list[FunctionSource], path: Path, original_code: str - ) -> tuple[str, dict[Path, str]]: - should_sort_imports = not self.args.disable_imports_sorting - if should_sort_imports and isort.code(original_code) != original_code: - should_sort_imports = False - - new_code = format_code(self.args.formatter_cmds, path) - if should_sort_imports: - new_code = sort_imports(new_code) - - new_helper_code: dict[Path, str] = {} - helper_functions_paths = {hf.file_path for hf in helper_functions} - for module_abspath in helper_functions_paths: - formatted_helper_code = format_code(self.args.formatter_cmds, module_abspath) - if should_sort_imports: - formatted_helper_code = sort_imports(formatted_helper_code) - new_helper_code[module_abspath] = formatted_helper_code - - return new_code, new_helper_code - - def replace_function_and_helpers_with_optimized_code( - self, - code_context: CodeOptimizationContext, - function_to_optimize_file_path: Path, - optimized_code: str, - qualified_function_name: str, - ) -> bool: - did_update = False - read_writable_functions_by_file_path = defaultdict(set) - read_writable_functions_by_file_path[function_to_optimize_file_path].add(qualified_function_name) - for helper_function in code_context.helper_functions: - if helper_function.jedi_definition.type != "class": - read_writable_functions_by_file_path[helper_function.file_path].add(helper_function.qualified_name) - for module_abspath, qualified_names in read_writable_functions_by_file_path.items(): - did_update |= replace_function_definitions_in_module( - function_names=list(qualified_names), - optimized_code=optimized_code, - module_abspath=module_abspath, - preexisting_objects=code_context.preexisting_objects, - project_root_path=self.args.project_root, - ) - return did_update - - def get_code_optimization_context( - self, function_to_optimize: FunctionToOptimize, project_root: Path, original_source_code: str - ) -> Result[CodeOptimizationContext, str]: - code_to_optimize, contextual_dunder_methods = extract_code([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( - function_to_optimize, self.args.project_root, code_to_optimize - ) - if function_to_optimize.parents: - function_class = 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 - ] + [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( - original_source_code, - code_to_optimize_with_helpers, - function_to_optimize.file_path, - function_to_optimize.file_path, - project_root, - helper_functions, - ) - - try: - new_code_ctx = code_context_extractor.get_code_optimization_context(function_to_optimize, project_root) - except ValueError as e: - return Failure(str(e)) - - return Success( - CodeOptimizationContext( - 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 - preexisting_objects=new_code_ctx.preexisting_objects, - ) - ) - - @staticmethod - def cleanup_leftover_test_return_values() -> None: - # remove leftovers from previous run - get_run_tmp_file(Path("test_return_values_0.bin")).unlink(missing_ok=True) - get_run_tmp_file(Path("test_return_values_0.sqlite")).unlink(missing_ok=True) - - def instrument_existing_tests( - self, function_to_optimize: FunctionToOptimize, function_to_tests: dict[str, list[FunctionCalledInTest]] - ) -> set[Path]: - existing_test_files_count = 0 - replay_test_files_count = 0 - concolic_coverage_test_files_count = 0 - unique_instrumented_test_files = set() - - func_qualname = function_to_optimize.qualified_name_with_modules_from_root(self.args.project_root) - if func_qualname not in function_to_tests: - logger.info(f"Did not find any pre-existing tests for '{func_qualname}', will only use generated tests.") - console.rule() - else: - test_file_invocation_positions = defaultdict(list[FunctionCalledInTest]) - for tests_in_file in function_to_tests.get(func_qualname): - test_file_invocation_positions[ - (tests_in_file.tests_in_file.test_file, tests_in_file.tests_in_file.test_type) - ].append(tests_in_file) - for (test_file, test_type), tests_in_file_list in test_file_invocation_positions.items(): - path_obj_test_file = Path(test_file) - if test_type == TestType.EXISTING_UNIT_TEST: - existing_test_files_count += 1 - elif test_type == TestType.REPLAY_TEST: - replay_test_files_count += 1 - elif test_type == TestType.CONCOLIC_COVERAGE_TEST: - concolic_coverage_test_files_count += 1 - else: - msg = f"Unexpected test type: {test_type}" - raise ValueError(msg) - success, injected_behavior_test = inject_profiling_into_existing_test( - mode=TestingMode.BEHAVIOR, - test_path=path_obj_test_file, - call_positions=[test.position for test in tests_in_file_list], - function_to_optimize=function_to_optimize, - tests_project_root=self.test_cfg.tests_project_rootdir, - test_framework=self.args.test_framework, - ) - if not success: - continue - success, injected_perf_test = inject_profiling_into_existing_test( - mode=TestingMode.PERFORMANCE, - test_path=path_obj_test_file, - call_positions=[test.position for test in tests_in_file_list], - function_to_optimize=function_to_optimize, - tests_project_root=self.test_cfg.tests_project_rootdir, - test_framework=self.args.test_framework, - ) - if not success: - continue - # TODO: this naming logic should be moved to a function and made more standard - new_behavioral_test_path = Path( - f"{os.path.splitext(test_file)[0]}__perfinstrumented{os.path.splitext(test_file)[1]}" - ) - new_perf_test_path = Path( - f"{os.path.splitext(test_file)[0]}__perfonlyinstrumented{os.path.splitext(test_file)[1]}" - ) - if injected_behavior_test is not None: - with new_behavioral_test_path.open("w", encoding="utf8") as _f: - _f.write(injected_behavior_test) - else: - msg = "injected_behavior_test is None" - raise ValueError(msg) - if injected_perf_test is not None: - with new_perf_test_path.open("w", encoding="utf8") as _f: - _f.write(injected_perf_test) - - 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( - instrumented_behavior_file_path=new_behavioral_test_path, - benchmarking_file_path=new_perf_test_path, - original_source=None, - original_file_path=Path(test_file), - test_type=test_type, - 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" - f"{'s' if replay_test_files_count != 1 else ''}, and " - f"{concolic_coverage_test_files_count} concolic coverage test file" - f"{'s' if concolic_coverage_test_files_count != 1 else ''} for {func_qualname}" - ) - return unique_instrumented_test_files - - def generate_tests_and_optimizations( - self, - code_to_optimize_with_helpers: str, - read_writable_code: str, - read_only_context_code: str, - function_to_optimize: FunctionToOptimize, - helper_functions: list[FunctionSource], - module_path: Path, - function_trace_id: str, - generated_test_paths: list[Path], - generated_perf_test_paths: list[Path], - function_to_optimize_ast: ast.FunctionDef, - run_experiment: bool = False, - ) -> Result[tuple[GeneratedTestsList, dict[str, list[FunctionCalledInTest]], OptimizationSet], str]: - assert len(generated_test_paths) == N_TESTS_TO_GENERATE - max_workers = N_TESTS_TO_GENERATE + 2 if not run_experiment else N_TESTS_TO_GENERATE + 3 - console.rule() - with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: - # Submit the test generation task as future - future_tests = self.generate_and_instrument_tests( - executor, - code_to_optimize_with_helpers, - function_to_optimize, - [definition.fully_qualified_name for definition in helper_functions], - module_path, - generated_test_paths, - generated_perf_test_paths, - (function_trace_id[:-4] + "EXP0" if run_experiment else function_trace_id), - ) - future_optimization_candidates = executor.submit( - self.aiservice_client.optimize_python_code, - read_writable_code, - read_only_context_code, - function_trace_id[:-4] + "EXP0" if run_experiment else function_trace_id, - N_CANDIDATES, - ExperimentMetadata(id=self.experiment_id, group="control") if run_experiment else None, - ) - future_candidates_exp = None - - future_concolic_tests = executor.submit( - generate_concolic_tests, self.test_cfg, self.args, function_to_optimize, function_to_optimize_ast - ) - futures = [*future_tests, future_optimization_candidates, future_concolic_tests] - if run_experiment: - future_candidates_exp = executor.submit( - self.local_aiservice_client.optimize_python_code, - read_writable_code, - read_only_context_code, - function_trace_id[:-4] + "EXP1", - N_CANDIDATES, - ExperimentMetadata(id=self.experiment_id, group="experiment"), - ) - futures.append(future_candidates_exp) - - # Wait for all futures to complete - concurrent.futures.wait(futures) - - # Retrieve results - candidates: list[OptimizedCandidate] = future_optimization_candidates.result() - if not candidates: - return Failure(f"/!\\ NO OPTIMIZATIONS GENERATED for {function_to_optimize.function_name}") - - candidates_experiment = future_candidates_exp.result() if future_candidates_exp else None - - # Process test generation results - - tests: list[GeneratedTests] = [] - for future in future_tests: - res = future.result() - if res: - ( - generated_test_source, - instrumented_behavior_test_source, - instrumented_perf_test_source, - test_behavior_path, - test_perf_path, - ) = res - tests.append( - GeneratedTests( - generated_original_test_source=generated_test_source, - instrumented_behavior_test_source=instrumented_behavior_test_source, - instrumented_perf_test_source=instrumented_perf_test_source, - behavior_file_path=test_behavior_path, - perf_file_path=test_perf_path, - ) - ) - if not tests: - logger.warning(f"Failed to generate and instrument tests for {function_to_optimize.function_name}") - return Failure(f"/!\\ NO TESTS GENERATED for {function_to_optimize.function_name}") - function_to_concolic_tests, concolic_test_str = future_concolic_tests.result() - logger.info(f"Generated {len(tests)} tests for {function_to_optimize.function_name}") - console.rule() - generated_tests = GeneratedTestsList(generated_tests=tests) - - return Success( - ( - generated_tests, - function_to_concolic_tests, - concolic_test_str, - OptimizationSet(control=candidates, experiment=candidates_experiment), - ) - ) - - def establish_original_code_baseline( - self, - function_name: str, - function_file_path: Path, - code_context: CodeOptimizationContext, - function_to_optimize: FunctionToOptimize, - original_helper_code: dict[Path, str], - file_path_to_helper_classes: dict[Path, set[str]], - ) -> Result[tuple[OriginalCodeBaseline, list[str]], str]: - # For the original function - run the tests and get the runtime, plus coverage - with progress_bar(f"Establishing original code baseline for {function_name}"): - assert (test_framework := self.args.test_framework) in ["pytest", "unittest"] - success = True - - test_env = os.environ.copy() - test_env["CODEFLASH_TEST_ITERATION"] = "0" - test_env["CODEFLASH_TRACER_DISABLE"] = "1" - if "PYTHONPATH" not in test_env: - test_env["PYTHONPATH"] = str(self.args.project_root) - else: - test_env["PYTHONPATH"] += os.pathsep + str(self.args.project_root) - - coverage_results = None - original_fto_code = function_file_path.read_text("utf-8") - # Instrument codeflash capture - try: - instrument_codeflash_capture( - function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root - ) - behavioral_results, coverage_results = self.run_and_parse_tests( - testing_type=TestingMode.BEHAVIOR, - test_env=test_env, - test_files=self.test_files, - optimization_iteration=0, - testing_time=TOTAL_LOOPING_TIME, - enable_coverage=test_framework == "pytest", - function_name=function_name, - source_file=function_file_path, - code_context=code_context, - ) - finally: - # Remove codeflash capture - self.write_code_and_helpers(original_fto_code, original_helper_code, function_to_optimize.file_path) - if test_framework == "pytest": - benchmarking_results, _ = self.run_and_parse_tests( - testing_type=TestingMode.PERFORMANCE, - test_env=test_env, - test_files=self.test_files, - optimization_iteration=0, - testing_time=TOTAL_LOOPING_TIME, - enable_coverage=False, - function_name=function_name, - source_file=function_file_path, - code_context=code_context, - ) - - else: - benchmarking_results = TestResults() - start_time: float = time.time() - for i in range(100): - if i >= 5 and time.time() - start_time >= TOTAL_LOOPING_TIME * 1.5: - # * 1.5 to give unittest a bit more time to run - break - test_env["CODEFLASH_LOOP_INDEX"] = str(i + 1) - unittest_loop_results, _ = self.run_and_parse_tests( - testing_type=TestingMode.PERFORMANCE, - test_env=test_env, - test_files=self.test_files, - optimization_iteration=0, - testing_time=TOTAL_LOOPING_TIME, - enable_coverage=False, - function_name=function_name, - source_file=function_file_path, - code_context=code_context, - unittest_loop_index=i + 1, - ) - benchmarking_results.merge(unittest_loop_results) - - console.print( - TestResults.report_to_tree( - behavioral_results.get_test_pass_fail_report_by_type(), - title="Overall test results for original code", - ) - ) - console.rule() - - total_timing = benchmarking_results.total_passed_runtime() # caution: doesn't handle the loop index - functions_to_remove = [ - result.id.test_function_name - for result in behavioral_results - if (result.test_type == TestType.GENERATED_REGRESSION and not result.did_pass) - ] - - if not behavioral_results: - logger.warning( - f"Couldn't run any tests for original function {function_name}. SKIPPING OPTIMIZING THIS FUNCTION." - ) - console.rule() - success = False - if total_timing == 0: - logger.warning( - "The overall summed benchmark runtime of the original function is 0, couldn't run tests." - ) - console.rule() - success = False - if not total_timing: - logger.warning("Failed to run the tests for the original function, skipping optimization") - console.rule() - success = False - if not success: - return Failure("Failed to establish a baseline for the original code.") - - loop_count = max([int(result.loop_index) for result in benchmarking_results.test_results]) - logger.info( - f"Original code summed runtime measured over {loop_count} loop{'s' if loop_count > 1 else ''}: " - f"{humanize_runtime(total_timing)} per full loop" - ) - console.rule() - logger.debug(f"Total original code runtime (ns): {total_timing}") - return Success( - ( - OriginalCodeBaseline( - behavioral_test_results=behavioral_results, - benchmarking_test_results=benchmarking_results, - runtime=total_timing, - coverage_results=coverage_results, - ), - functions_to_remove, - ) - ) - - def run_optimized_candidate( - self, - *, - optimization_candidate_index: int, - baseline_results: OriginalCodeBaseline, - function_to_optimize: FunctionToOptimize, - original_helper_code: dict[Path, str], - 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_TEST_ITERATION"] = str(optimization_candidate_index) - test_env["CODEFLASH_TRACER_DISABLE"] = "1" - if "PYTHONPATH" not in test_env: - test_env["PYTHONPATH"] = str(self.args.project_root) - else: - test_env["PYTHONPATH"] += os.pathsep + str(self.args.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(function_to_optimize.file_path).read_text("utf-8") - candidate_helper_code = {} - for module_abspath in original_helper_code: - candidate_helper_code[module_abspath] = Path(module_abspath).read_text("utf-8") - try: - instrument_codeflash_capture( - function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root - ) - - candidate_behavior_results, _ = self.run_and_parse_tests( - testing_type=TestingMode.BEHAVIOR, - test_env=test_env, - test_files=self.test_files, - optimization_iteration=optimization_candidate_index, - testing_time=TOTAL_LOOPING_TIME, - enable_coverage=False, - ) - # Remove instrumentation - finally: - self.write_code_and_helpers(candidate_fto_code, candidate_helper_code, function_to_optimize.file_path) - console.print( - TestResults.report_to_tree( - candidate_behavior_results.get_test_pass_fail_report_by_type(), - title="Behavioral Test Results for candidate", - ) - ) - console.rule() - - if compare_test_results(baseline_results.behavioral_test_results, candidate_behavior_results): - logger.info("Test results matched!") - console.rule() - else: - logger.info("Test results did not match the test results of the original code.") - console.rule() - return Failure("Test results did not match the test results of the original code.") - - if test_framework == "pytest": - candidate_benchmarking_results, _ = self.run_and_parse_tests( - testing_type=TestingMode.PERFORMANCE, - test_env=test_env, - test_files=self.test_files, - optimization_iteration=optimization_candidate_index, - testing_time=TOTAL_LOOPING_TIME, - enable_coverage=False, - ) - loop_count = ( - max(all_loop_indices) - if ( - all_loop_indices := { - result.loop_index for result in candidate_benchmarking_results.test_results - } - ) - else 0 - ) - - else: - candidate_benchmarking_results = TestResults() - start_time: float = time.time() - loop_count = 0 - for i in range(100): - if i >= 5 and time.time() - start_time >= TOTAL_LOOPING_TIME * 1.5: - # * 1.5 to give unittest a bit more time to run - break - test_env["CODEFLASH_LOOP_INDEX"] = str(i + 1) - unittest_loop_results, cov = self.run_and_parse_tests( - testing_type=TestingMode.PERFORMANCE, - test_env=test_env, - test_files=self.test_files, - optimization_iteration=optimization_candidate_index, - testing_time=TOTAL_LOOPING_TIME, - unittest_loop_index=i + 1, - ) - loop_count = i + 1 - candidate_benchmarking_results.merge(unittest_loop_results) - - if (total_candidate_timing := candidate_benchmarking_results.total_passed_runtime()) == 0: - logger.warning("The overall test runtime of the optimized function is 0, couldn't run tests.") - console.rule() - - logger.debug(f"Total optimized code {optimization_candidate_index} runtime (ns): {total_candidate_timing}") - return Success( - OptimizedCandidateResult( - max_loop_count=loop_count, - best_test_runtime=total_candidate_timing, - behavior_test_results=candidate_behavior_results, - benchmarking_test_results=candidate_benchmarking_results, - optimization_candidate_index=optimization_candidate_index, - total_candidate_timing=total_candidate_timing, - ) - ) - - def run_and_parse_tests( - self, - testing_type: TestingMode, - test_env: dict[str, str], - test_files: TestFiles, - optimization_iteration: int, - testing_time: float = TOTAL_LOOPING_TIME, - *, - enable_coverage: bool = False, - pytest_min_loops: int = 5, - pytest_max_loops: int = 100_000, - function_name: str | None = None, - source_file: Path | None = None, - code_context: CodeOptimizationContext | None = None, - unittest_loop_index: int | None = None, - ) -> tuple[TestResults, CoverageData | None]: - coverage_database_file = None - try: - if testing_type == TestingMode.BEHAVIOR: - result_file_path, run_result, coverage_database_file = run_behavioral_tests( - test_files, - test_framework=self.args.test_framework, - cwd=self.args.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.PERFORMANCE: - result_file_path, run_result = run_benchmarking_tests( - test_files, - cwd=self.args.project_root, - test_env=test_env, - 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.args.test_framework, - ) - else: - 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" - ) - return TestResults(), None - if run_result.returncode != 0: - 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"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=function_name, - source_file=source_file, - code_context=code_context, - coverage_database_file=coverage_database_file, - ) - return results, coverage_results - - def generate_and_instrument_tests( - self, - executor: concurrent.futures.ThreadPoolExecutor, - source_code_being_tested: str, - function_to_optimize: FunctionToOptimize, - helper_function_names: list[str], - module_path: Path, - generated_test_paths: list[Path], - generated_perf_test_paths: list[Path], - function_trace_id: str, - ) -> list[concurrent.futures.Future]: - return [ - executor.submit( - generate_tests, - self.aiservice_client, - source_code_being_tested, - function_to_optimize, - helper_function_names, - module_path, - self.test_cfg, - INDIVIDUAL_TESTCASE_TIMEOUT, - function_trace_id, - test_index, - test_path, - test_perf_path, - ) - for test_index, (test_path, test_perf_path) in enumerate( - zip(generated_test_paths, generated_perf_test_paths) - ) - ] def run_with_args(args: Namespace) -> None: diff --git a/codeflash/verification/test_results.py b/codeflash/verification/test_results.py index 7491ce603..28d8bfc0d 100644 --- a/codeflash/verification/test_results.py +++ b/codeflash/verification/test_results.py @@ -168,7 +168,7 @@ 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}" + 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/codeflash/verification/test_runner.py b/codeflash/verification/test_runner.py index 7df601922..46203e65a 100644 --- a/codeflash/verification/test_runner.py +++ b/codeflash/verification/test_runner.py @@ -87,7 +87,8 @@ def run_behavioral_tests( env=pytest_test_env, timeout=600, ) - logger.debug(results) + logger.debug( + f"""Result return code: {results.returncode}, {"Result stderr:" + str(results.stderr) if results.stderr else ''}""") else: results = execute_test_subprocess( pytest_cmd_list + common_pytest_args + result_args + test_files, @@ -95,12 +96,16 @@ def run_behavioral_tests( env=pytest_test_env, timeout=600, # TODO: Make this dynamic ) + logger.debug( + f"""Result return code: {results.returncode}, {"Result stderr:" + str(results.stderr) if results.stderr else ''}""") elif test_framework == "unittest": if enable_coverage: raise ValueError("Coverage is not supported yet for unittest framework") test_env["CODEFLASH_LOOP_INDEX"] = "1" test_files = [file.instrumented_behavior_file_path for file in test_paths.test_files] result_file_path, results = run_unittest_tests(verbose, test_files, test_env, cwd) + logger.debug( + f"""Result return code: {results.returncode}, {"Result stderr:" + str(results.stderr) if results.stderr else ''}""") else: raise ValueError(f"Unsupported test framework: {test_framework}") diff --git a/pyproject.toml b/pyproject.toml index a7c1696c4..62587fa9d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,6 +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" [tool.poetry.build] script = "codeflash/update_license_version.py" diff --git a/tests/scripts/end_to_end_test_init_optimization.py b/tests/scripts/end_to_end_test_init_optimization.py index 4f4855a1e..a19be5d82 100644 --- a/tests/scripts/end_to_end_test_init_optimization.py +++ b/tests/scripts/end_to_end_test_init_optimization.py @@ -1,7 +1,7 @@ import os import pathlib -from end_to_end_test_utilities import TestConfig, run_codeflash_command, run_with_retries +from end_to_end_test_utilities import CoverageExpectation, TestConfig, run_codeflash_command, run_with_retries def run_test(expected_improvement_pct: int) -> bool: @@ -10,6 +10,11 @@ def run_test(expected_improvement_pct: int) -> bool: function_name="CharacterRemover.remove_control_characters", test_framework="pytest", min_improvement_x=1.0, + coverage_expectations=[ + CoverageExpectation( + function_name="CharacterRemover.remove_control_characters", expected_coverage=100.0, expected_lines=[14] + ) + ], ) cwd = (pathlib.Path(__file__).parent.parent.parent / "code_to_optimize").resolve() return run_codeflash_command(cwd, config, expected_improvement_pct) diff --git a/tests/test_code_replacement.py b/tests/test_code_replacement.py index f1801fe20..501a0583b 100644 --- a/tests/test_code_replacement.py +++ b/tests/test_code_replacement.py @@ -2,7 +2,6 @@ import dataclasses import os -from argparse import Namespace from collections import defaultdict from pathlib import Path @@ -14,7 +13,8 @@ ) from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.models.models import FunctionParent -from codeflash.optimization.optimizer import Optimizer +from codeflash.optimization.function_optimizer import FunctionOptimizer +from codeflash.verification.verification_utils import TestConfig os.environ["CODEFLASH_API_KEY"] = "cf-test-key" @@ -766,24 +766,18 @@ def main_method(self): return HelperClass(self.name).helper_method() """ file_path = Path(__file__).resolve() - opt = Optimizer( - Namespace( - project_root=file_path.parent.resolve(), - disable_telemetry=True, - tests_root="tests", - test_framework="pytest", - pytest_cmd="pytest", - experiment_id=None, - test_project_root=file_path.parent.resolve(), - ) - ) func_top_optimize = FunctionToOptimize( function_name="main_method", file_path=file_path, parents=[FunctionParent("MainClass", "ClassDef")] ) - original_code = file_path.read_text() - code_context = opt.get_code_optimization_context( - function_to_optimize=func_top_optimize, project_root=file_path.parent, original_source_code=original_code - ).unwrap() + test_config = TestConfig( + tests_root=file_path.parent, + tests_project_rootdir=file_path.parent, + project_root_path=file_path.parent, + test_framework="pytest", + pytest_cmd="pytest", + ) + func_optimizer = FunctionOptimizer(function_to_optimize=func_top_optimize, test_cfg=test_config) + code_context = func_optimizer.get_code_optimization_context().unwrap() assert code_context.code_to_optimize_with_helpers == get_code_output @@ -1013,35 +1007,35 @@ def to_name(self) -> str: class TestResults(BaseModel): 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 __delitem__(self, index: int) -> None: del self.test_results[index] - + 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 not isinstance(other, TestResults) or len(self) != len(other): return False - + # Increase recursion limit only if necessary original_recursion_limit = sys.getrecursionlimit() if original_recursion_limit < 5000: sys.setrecursionlimit(5000) - + for test_result in self: other_test_result = other.get_by_id(test_result.id) if other_test_result is None or not ( @@ -1054,10 +1048,10 @@ def __eq__(self, other: object) -> bool: ): sys.setrecursionlimit(original_recursion_limit) return False - + sys.setrecursionlimit(original_recursion_limit) return True - + def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]: report = {test_type: {"passed": 0, "failed": 0} for test_type in TestType} for test_result in self.test_results: @@ -1105,8 +1099,8 @@ def get_test_pass_fail_report_by_type(self) -> dict[TestType, dict[str, int]]: ) assert ( - new_code - == """from __future__ import annotations + new_code + == """from __future__ import annotations import sys from codeflash.verification.comparator import comparator from enum import Enum @@ -1245,21 +1239,21 @@ def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray: """Row-wise cosine similarity between two equal-width matrices.""" if len(X.data) == 0 or len(Y.data) == 0: return np.array([]) - + X_np, Y_np = np.asarray(X.data), np.asarray(Y.data) if X_np.shape[1] != Y_np.shape[1]: raise ValueError(f"Number of columns in X and Y must be the same. X has shape {X_np.shape} and Y has shape {Y_np.shape}.") X_norm = np.linalg.norm(X_np, axis=1, keepdims=True) Y_norm = np.linalg.norm(Y_np, axis=1, keepdims=True) - + norm_product = X_norm * Y_norm.T norm_product[norm_product == 0] = np.inf # Prevent division by zero dot_product = np.dot(X_np, Y_np.T) similarity = dot_product / norm_product - + # Any NaN or Inf values are set to 0.0 np.nan_to_num(similarity, copy=False) - + return similarity def cosine_similarity_top_k( X: Matrix, @@ -1270,15 +1264,15 @@ def cosine_similarity_top_k( """Row-wise cosine similarity with optional top-k and score threshold filtering.""" if len(X.data) == 0 or len(Y.data) == 0: return [], [] - + score_array = cosine_similarity(X, Y) - + sorted_idxs = np.argpartition(-score_array.flatten(), range(top_k or len(score_array.flatten())))[:(top_k or len(score_array.flatten()))] sorted_idxs = sorted_idxs[score_array.flatten()[sorted_idxs] > (score_threshold if score_threshold is not None else -1)] - + ret_idxs = [(x // score_array.shape[1], x % score_array.shape[1]) for x in sorted_idxs] scores = score_array.flatten()[sorted_idxs].tolist() - + return ret_idxs, scores ''' preexisting_objects: list[tuple[str, list[FunctionParent]]] = find_preexisting_objects(original_code) @@ -1311,8 +1305,8 @@ def cosine_similarity_top_k( project_root_path=Path(__file__).parent.parent.resolve(), ) assert ( - new_code - == '''import numpy as np + new_code + == '''import numpy as np from pydantic.dataclasses import dataclass from typing import List, Optional, Tuple, Union @dataclass(config=dict(arbitrary_types_allowed=True)) @@ -1343,15 +1337,15 @@ def cosine_similarity_top_k( """Row-wise cosine similarity with optional top-k and score threshold filtering.""" if len(X.data) == 0 or len(Y.data) == 0: return [], [] - + score_array = cosine_similarity(X, Y) - + sorted_idxs = np.argpartition(-score_array.flatten(), range(top_k or len(score_array.flatten())))[:(top_k or len(score_array.flatten()))] sorted_idxs = sorted_idxs[score_array.flatten()[sorted_idxs] > (score_threshold if score_threshold is not None else -1)] - + ret_idxs = [(x // score_array.shape[1], x % score_array.shape[1]) for x in sorted_idxs] scores = score_array.flatten()[sorted_idxs].tolist() - + return ret_idxs, scores ''' ) @@ -1370,8 +1364,8 @@ def cosine_similarity_top_k( ) assert ( - new_helper_code - == '''import numpy as np + new_helper_code + == '''import numpy as np from pydantic.dataclasses import dataclass from typing import List, Optional, Tuple, Union @dataclass(config=dict(arbitrary_types_allowed=True)) @@ -1381,21 +1375,21 @@ def cosine_similarity(X: Matrix, Y: Matrix) -> np.ndarray: """Row-wise cosine similarity between two equal-width matrices.""" if len(X.data) == 0 or len(Y.data) == 0: return np.array([]) - + X_np, Y_np = np.asarray(X.data), np.asarray(Y.data) if X_np.shape[1] != Y_np.shape[1]: raise ValueError(f"Number of columns in X and Y must be the same. X has shape {X_np.shape} and Y has shape {Y_np.shape}.") X_norm = np.linalg.norm(X_np, axis=1, keepdims=True) Y_norm = np.linalg.norm(Y_np, axis=1, keepdims=True) - + norm_product = X_norm * Y_norm.T norm_product[norm_product == 0] = np.inf # Prevent division by zero dot_product = np.dot(X_np, Y_np.T) similarity = dot_product / norm_product - + # Any NaN or Inf values are set to 0.0 np.nan_to_num(similarity, copy=False) - + return similarity def cosine_similarity_top_k( X: Matrix, @@ -1406,15 +1400,15 @@ def cosine_similarity_top_k( """Row-wise cosine similarity with optional top-k and score threshold filtering.""" if len(X.data) == 0 or len(Y.data) == 0: return [], [] - + score_array = cosine_similarity(X, Y) - + sorted_idxs = np.argpartition(-score_array.flatten(), range(top_k or len(score_array.flatten())))[:(top_k or len(score_array.flatten()))] sorted_idxs = sorted_idxs[score_array.flatten()[sorted_idxs] > (score_threshold if score_threshold is not None else -1)] - + ret_idxs = [(x // score_array.shape[1], x % score_array.shape[1]) for x in sorted_idxs] scores = score_array.flatten()[sorted_idxs].tolist() - + return ret_idxs, scores ''' ) @@ -1481,7 +1475,7 @@ def test_future_aliased_imports_removal() -> None: def test_0_diff_code_replacement(): original_code = """from __future__ import annotations - + import numpy as np def functionA(): return np.array([1, 2, 3]) diff --git a/tests/test_codeflash_capture.py b/tests/test_codeflash_capture.py index 1e8426aef..83b1efd2b 100644 --- a/tests/test_codeflash_capture.py +++ b/tests/test_codeflash_capture.py @@ -2,18 +2,18 @@ import os import re -from argparse import Namespace from pathlib import Path from codeflash.code_utils.code_utils import get_run_tmp_file from codeflash.code_utils.compat import SAFE_SYS_EXECUTABLE from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.models.models import FunctionParent, TestFile, TestFiles, TestingMode -from codeflash.optimization.optimizer import Optimizer +from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.verification.equivalence import compare_test_results from codeflash.verification.instrument_codeflash_capture import instrument_codeflash_capture from codeflash.verification.test_results import TestType, VerificationType from codeflash.verification.test_runner import execute_test_subprocess +from codeflash.verification.verification_utils import TestConfig # Tests for get_stack_info. Ensures that when a test is run via pytest, the correct test information is extracted @@ -184,6 +184,7 @@ def __init__(self): self.x = 2 print(f"TEST_INFO_START|{{get_test_info_from_stack('{test_dir!s}')}}|TEST_INFO_END") """ + test_dir = (Path(__file__).parent.parent / "code_to_optimize" / "tests" / "pytest").resolve() test_file_name = "test_stack_info_temp.py" test_path = test_dir / test_file_name @@ -319,7 +320,6 @@ def __init__(self): assert results[5][3] == "23" finally: - # Clean up files test_path.unlink(missing_ok=True) sample_code_path.unlink(missing_ok=True) @@ -346,6 +346,7 @@ def __init__(self): self.x = 2 print(f"TEST_INFO_START|{{get_test_info_from_stack('{test_dir!s}')}}|TEST_INFO_END") """ + test_dir = (Path(__file__).parent.parent / "code_to_optimize" / "tests" / "pytest").resolve() test_file_name = "test_stack_info_temp.py" test_path = test_dir / test_file_name @@ -431,23 +432,25 @@ def __init__(self, x=2): f.write(test_code) with sample_code_path.open("w") as f: f.write(sample_code) - 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 - test_files = TestFiles( + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="pytest", + pytest_cmd="pytest", + ) + fto = FunctionToOptimize( + function_name="some_function", + file_path=sample_code_path, + parents=[FunctionParent(name="MyClass", type="ClassDef")], + ) + func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config) + func_optimizer.test_files = TestFiles( test_files=[ TestFile( instrumented_behavior_file_path=test_path, @@ -457,10 +460,10 @@ def __init__(self, x=2): ) ] ) - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, @@ -541,23 +544,24 @@ def __init__(self, *args, **kwargs): f.write(test_code) with sample_code_path.open("w") as f: f.write(sample_code) - 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 - test_files = TestFiles( + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="pytest", + pytest_cmd="pytest", + ) + fto = FunctionToOptimize( + function_name="some_function", + file_path=sample_code_path, + parents=[FunctionParent(name="MyClass", type="ClassDef")], + ) + func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config) + func_optimizer.test_files = TestFiles( test_files=[ TestFile( instrumented_behavior_file_path=test_path, @@ -567,10 +571,10 @@ def __init__(self, *args, **kwargs): ) ] ) - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, @@ -653,25 +657,24 @@ def __init__(self, x=2): with sample_code_path.open("w") as f: f.write(sample_code) - 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, - ) - ) - - # Environment variables for codeflash test_env = os.environ.copy() test_env["CODEFLASH_TEST_ITERATION"] = "0" test_env["CODEFLASH_LOOP_INDEX"] = "1" - test_type = TestType.EXISTING_UNIT_TEST - test_files = TestFiles( + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="pytest", + pytest_cmd="pytest", + ) + fto = FunctionToOptimize( + function_name="some_function", + file_path=sample_code_path, + parents=[FunctionParent(name="MyClass", type="ClassDef")], + ) + func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config) + func_optimizer.test_files = TestFiles( test_files=[ TestFile( instrumented_behavior_file_path=test_path, @@ -681,12 +684,10 @@ def __init__(self, x=2): ) ] ) - - # Run the tests and parse results - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, @@ -804,24 +805,26 @@ def another_helper(self): f.write(original_code) with test_path.open("w") as f: f.write(test_code) - 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 - test_files = TestFiles( + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="pytest", + pytest_cmd="pytest", + ) + fto = FunctionToOptimize( + function_name="target_function", + file_path=fto_file_path, + parents=[FunctionParent(name="MyClass", type="ClassDef")], + ) + func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config) + func_optimizer.test_files = TestFiles( test_files=[ TestFile( instrumented_behavior_file_path=test_path, @@ -832,10 +835,10 @@ def another_helper(self): ] ) - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, @@ -925,7 +928,7 @@ def another_helper(self): helper_path_2 = test_dir / helper_file_2 fto_file_path = test_dir / fto_file_name - tests_root = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/").resolve() + tests_root = Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/" project_root_path = (Path(__file__).parent / "..").resolve() try: @@ -938,30 +941,26 @@ def another_helper(self): with test_path.open("w") as f: f.write(test_code) - fto = FunctionToOptimize("target_function", str(fto_file_path), parents=[FunctionParent("MyClass", "ClassDef")]) + fto = FunctionToOptimize("target_function", fto_file_path, parents=[FunctionParent("MyClass", "ClassDef")]) file_path_to_helper_class = { helper_path_1: {"HelperClass1"}, helper_path_2: {"HelperClass2", "AnotherHelperClass"}, } instrument_codeflash_capture(fto, file_path_to_helper_class, tests_root) - 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 - test_files = TestFiles( + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="pytest", + pytest_cmd="pytest", + ) + func_optimizer = FunctionOptimizer(function_to_optimize=fto, test_cfg=test_config) + func_optimizer.test_files = TestFiles( test_files=[ TestFile( instrumented_behavior_file_path=test_path, @@ -983,10 +982,10 @@ def another_helper(self): } instrument_codeflash_capture(fto, file_path_to_helper_classes, tests_root) - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, @@ -994,7 +993,7 @@ def another_helper(self): ) # Remove instrumentation - opt.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path) + FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path) assert len(test_results.test_results) == 4 assert test_results[0].id.test_function_name == "test_helper_classes" @@ -1035,17 +1034,17 @@ def target_function(self): Path(helper_path_2): {"HelperClass2", "AnotherHelperClass"}, } instrument_codeflash_capture(fto, file_path_to_helper_classes, tests_root) - modified_test_results, coverage_data = opt.run_and_parse_tests( + modified_test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, testing_time=0.1, ) # Remove instrumentation - opt.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path) + FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path) # Now, this fto_code mutates the instance so it should fail mutated_fto_code = """ @@ -1074,17 +1073,17 @@ def target_function(self): Path(helper_path_2): {"HelperClass2", "AnotherHelperClass"}, } instrument_codeflash_capture(fto, file_path_to_helper_classes, tests_root) - mutated_test_results, coverage_data = opt.run_and_parse_tests( + mutated_test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, testing_time=0.1, ) # Remove instrumentation - opt.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path) + FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path) assert not compare_test_results(test_results, mutated_test_results) # This fto code stopped using a helper class. it should still pass @@ -1112,17 +1111,17 @@ def target_function(self): Path(helper_path_2): {"HelperClass2", "AnotherHelperClass"}, } instrument_codeflash_capture(fto, file_path_to_helper_classes, tests_root) - no_helper1_test_results, coverage_data = opt.run_and_parse_tests( + no_helper1_test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, testing_time=0.1, ) # Remove instrumentation - opt.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path) + FunctionOptimizer.write_code_and_helpers(candidate_fto_code, candidate_helper_code, fto.file_path) assert compare_test_results(test_results, no_helper1_test_results) finally: diff --git a/tests/test_function_dependencies.py b/tests/test_function_dependencies.py index f45309087..faa754af9 100644 --- a/tests/test_function_dependencies.py +++ b/tests/test_function_dependencies.py @@ -1,5 +1,4 @@ import pathlib -from argparse import Namespace from dataclasses import dataclass import pytest @@ -7,7 +6,8 @@ from codeflash.either import is_successful from codeflash.models.models import FunctionParent from codeflash.optimization.function_context import get_function_variables_definitions -from codeflash.optimization.optimizer import Optimizer +from codeflash.optimization.function_optimizer import FunctionOptimizer +from codeflash.verification.verification_utils import TestConfig def calculate_something(data): @@ -184,17 +184,7 @@ def topologicalSort(self): def test_class_method_dependencies() -> None: file_path = pathlib.Path(__file__).resolve() - opt = Optimizer( - Namespace( - project_root=file_path.parent.resolve(), - disable_telemetry=True, - tests_root="tests", - test_framework="pytest", - pytest_cmd="pytest", - experiment_id=None, - test_project_root=file_path.parent.resolve(), - ) - ) + function_to_optimize = FunctionToOptimize( function_name="topologicalSort", file_path=str(file_path), @@ -202,9 +192,19 @@ def test_class_method_dependencies() -> None: starting_line=None, ending_line=None, ) + func_optimizer = FunctionOptimizer( + function_to_optimize=function_to_optimize, + test_cfg=TestConfig( + tests_root=file_path, + tests_project_rootdir=file_path.parent, + project_root_path=file_path.parent, + test_framework="pytest", + pytest_cmd="pytest", + ), + ) with open(file_path) as f: original_code = f.read() - ctx_result = opt.get_code_optimization_context(function_to_optimize, opt.args.project_root, original_code) + ctx_result = func_optimizer.get_code_optimization_context() if not is_successful(ctx_result): pytest.fail() code_context = ctx_result.unwrap() @@ -280,17 +280,7 @@ def test_decorator_dependencies() -> None: def test_recursive_function_context() -> None: file_path = pathlib.Path(__file__).resolve() - opt = Optimizer( - Namespace( - project_root=file_path.parent.resolve(), - disable_telemetry=True, - tests_root="tests", - test_framework="pytest", - pytest_cmd="pytest", - experiment_id=None, - test_project_root=file_path.parent.resolve(), - ) - ) + function_to_optimize = FunctionToOptimize( function_name="recursive", file_path=str(file_path), @@ -298,9 +288,20 @@ def test_recursive_function_context() -> None: starting_line=None, ending_line=None, ) + func_optimizer = FunctionOptimizer( + function_to_optimize=function_to_optimize, + test_cfg=TestConfig( + tests_root=file_path, + tests_project_rootdir=file_path.parent, + project_root_path=file_path.parent, + test_framework="pytest", + pytest_cmd="pytest", + ), + ) with open(file_path) as f: original_code = f.read() - ctx_result = opt.get_code_optimization_context(function_to_optimize, opt.args.project_root, original_code) + + ctx_result = func_optimizer.get_code_optimization_context() if not is_successful(ctx_result): pytest.fail() code_context = ctx_result.unwrap() diff --git a/tests/test_get_helper_code.py b/tests/test_get_helper_code.py index 2970dc89e..b7dde84a4 100644 --- a/tests/test_get_helper_code.py +++ b/tests/test_get_helper_code.py @@ -3,10 +3,13 @@ from pathlib import Path import pytest + from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.either import is_successful from codeflash.models.models import FunctionParent +from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.optimization.optimizer import Optimizer +from codeflash.verification.verification_utils import TestConfig class HelperClass: @@ -31,6 +34,7 @@ def test_get_outside_method_helper() -> None: experiment_id=None, ) ) + function_to_optimize = FunctionToOptimize( function_name="OptimizeMe", file_path=file_path, parents=[], starting_line=None, ending_line=None ) @@ -51,7 +55,7 @@ def test_flavio_typed_code_helper() -> None: _STORE_T = TypeVar("_STORE_T") class AbstractCacheBackend(CacheBackend, Protocol[_KEY_T, _STORE_T]): """Interface for cache backends used by the persistent cache decorator.""" - + def __init__(self) -> None: ... def hash_key( @@ -213,17 +217,6 @@ def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R: f.write(code) f.flush() file_path = Path(f.name).resolve() - opt = Optimizer( - Namespace( - project_root=file_path.parent.resolve(), - disable_telemetry=True, - tests_root="tests", - test_framework="pytest", - pytest_cmd="pytest", - experiment_id=None, - test_project_root=Path().resolve(), - ) - ) function_to_optimize = FunctionToOptimize( function_name="__call__", file_path=file_path, @@ -231,17 +224,25 @@ def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R: starting_line=None, ending_line=None, ) + test_config = TestConfig( + tests_root="tests", + tests_project_rootdir=Path.cwd(), + project_root_path=file_path.parent.resolve(), + test_framework="pytest", + pytest_cmd="pytest", + ) + func_optimizer = FunctionOptimizer(function_to_optimize=function_to_optimize, test_cfg=test_config) with open(file_path) as f: original_code = f.read() - ctx_result = opt.get_code_optimization_context(function_to_optimize, opt.args.project_root, original_code) + ctx_result = func_optimizer.get_code_optimization_context() if not is_successful(ctx_result): pytest.fail() code_context = ctx_result.unwrap() assert code_context.helper_functions[0].qualified_name == "AbstractCacheBackend.get_cache_or_call" assert ( - code_context.code_to_optimize_with_helpers - == '''_R = TypeVar("_R") + code_context.code_to_optimize_with_helpers + == '''_R = TypeVar("_R") class AbstractCacheBackend(CacheBackend, Protocol[_KEY_T, _STORE_T]): def __init__(self) -> None: ... @@ -338,29 +339,27 @@ def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R: def test_bubble_sort_deps() -> None: file_path = (Path(__file__) / ".." / ".." / "code_to_optimize" / "bubble_sort_deps.py").resolve() - opt = Optimizer( - Namespace( - project_root=file_path.parent.parent.resolve(), - disable_telemetry=True, - tests_root=str(file_path.parent / "tests"), - test_framework="pytest", - pytest_cmd="pytest", - experiment_id=None, - test_project_root=file_path.parent.resolve(), - ) - ) + function_to_optimize = FunctionToOptimize( function_name="sorter_deps", file_path=file_path, parents=[], starting_line=None, ending_line=None ) + test_config = TestConfig( + tests_root=str(file_path.parent / "tests"), + tests_project_rootdir=file_path.parent.resolve(), + project_root_path=file_path.parent.parent.resolve(), + test_framework="pytest", + pytest_cmd="pytest", + ) + func_optimizer = FunctionOptimizer(function_to_optimize=function_to_optimize, test_cfg=test_config) with open(file_path) as f: original_code = f.read() - ctx_result = opt.get_code_optimization_context(function_to_optimize, opt.args.project_root, original_code) + ctx_result = func_optimizer.get_code_optimization_context() if not is_successful(ctx_result): pytest.fail() code_context = ctx_result.unwrap() assert ( - code_context.code_to_optimize_with_helpers - == """def dep1_comparer(arr, j: int) -> bool: + code_context.code_to_optimize_with_helpers + == """def dep1_comparer(arr, j: int) -> bool: return arr[j] > arr[j + 1] def dep2_swap(arr, j): @@ -378,7 +377,7 @@ def sorter_deps(arr): ) assert len(code_context.helper_functions) == 2 assert ( - code_context.helper_functions[0].fully_qualified_name - == "code_to_optimize.bubble_sort_dep1_helper.dep1_comparer" + code_context.helper_functions[0].fully_qualified_name + == "code_to_optimize.bubble_sort_dep1_helper.dep1_comparer" ) assert code_context.helper_functions[1].fully_qualified_name == "code_to_optimize.bubble_sort_dep2_swap.dep2_swap" diff --git a/tests/test_get_read_only_code.py b/tests/test_get_read_only_code.py index dd10eb5bc..0c71d9b6c 100644 --- a/tests/test_get_read_only_code.py +++ b/tests/test_get_read_only_code.py @@ -1,6 +1,7 @@ from textwrap import dedent import pytest + from codeflash.context.code_context_extractor import get_read_only_code @@ -66,7 +67,7 @@ def target_method(self): expected = """ class TestClass: - + def __str__(self): return f"Value: {self.x}" """ @@ -91,7 +92,7 @@ def target_method(self): expected = """ class TestClass: - + def __str__(self): return f"Value: {self.x}" """ @@ -118,7 +119,7 @@ def target_method(self): expected = """ class TestClass: - + def __str__(self): return f"Value: {self.x}" """ @@ -655,7 +656,7 @@ def __repr__(self) -> str: processor = DataProcessor(sample_data) class ResultHandler: - + def __str__(self) -> str: return f"ResultHandler(cache_size={len(self.cache)})" diff --git a/tests/test_instrument_all_and_run.py b/tests/test_instrument_all_and_run.py index c5005a575..643d4bde7 100644 --- a/tests/test_instrument_all_and_run.py +++ b/tests/test_instrument_all_and_run.py @@ -46,7 +46,7 @@ """ -def test_function_full_instrumentation() -> None: +def test_bubble_sort_behavior_results() -> None: code = """from code_to_optimize.bubble_sort import sorter @@ -147,7 +147,9 @@ def test_sort(): test_env["CODEFLASH_TEST_ITERATION"] = "0" test_env["CODEFLASH_LOOP_INDEX"] = "1" test_type = TestType.EXISTING_UNIT_TEST - test_files = TestFiles( + + func_optimizer = opt.create_function_optimizer(func) + func_optimizer.test_files = TestFiles( test_files=[ TestFile( instrumented_behavior_file_path=test_path, @@ -157,10 +159,10 @@ def test_sort(): ) ] ) - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, @@ -318,7 +320,8 @@ def test_sort(): test_env["CODEFLASH_TEST_ITERATION"] = "0" test_env["CODEFLASH_LOOP_INDEX"] = "1" test_type = TestType.EXISTING_UNIT_TEST - test_files = TestFiles( + func_optimizer = opt.create_function_optimizer(fto) + func_optimizer.test_files = TestFiles( test_files=[ TestFile( instrumented_behavior_file_path=test_path, @@ -328,10 +331,10 @@ def test_sort(): ) ] ) - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, @@ -412,11 +415,21 @@ def sorter(self, arr): test_project_root=project_root_path, ) ) - - new_test_results, coverage_data = opt.run_and_parse_tests( + func_optimizer = opt.create_function_optimizer(fto) + 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, + ) + ] + ) + new_test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, - test_files=test_files, + test_files=func_optimizer.test_files, optimization_iteration=0, pytest_min_loops=1, pytest_max_loops=1, diff --git a/tests/test_instrument_tests.py b/tests/test_instrument_tests.py index 49439c063..bf7373522 100644 --- a/tests/test_instrument_tests.py +++ b/tests/test_instrument_tests.py @@ -5,7 +5,6 @@ import os import sys import tempfile -from argparse import Namespace from pathlib import Path from codeflash.code_utils.code_utils import get_run_tmp_file @@ -15,8 +14,9 @@ ) from codeflash.discovery.functions_to_optimize import FunctionToOptimize from codeflash.models.models import CodePosition, FunctionParent, TestFile, TestFiles, TestingMode, TestsInFile -from codeflash.optimization.optimizer import Optimizer +from codeflash.optimization.function_optimizer import FunctionOptimizer from codeflash.verification.test_results import TestType +from codeflash.verification.verification_utils import TestConfig codeflash_wrap_string = """def codeflash_wrap(wrapped, test_module_name, test_class_name, test_name, function_name, line_id, loop_index, codeflash_cur, codeflash_con, *args, **kwargs): test_id = f'{{test_module_name}}:{{test_class_name}}:{{test_name}}:{{line_id}}:{{loop_index}}' @@ -160,7 +160,7 @@ def test_sort(self): with tempfile.NamedTemporaryFile(mode="w") as f: f.write(code) f.flush() - func = FunctionToOptimize(function_name="sorter", parents=[], file_path=Path("module.py")) + func = FunctionToOptimize(function_name="sorter", parents=[], file_path=Path(f.name)) original_cwd = Path.cwd() run_cwd = Path(__file__).parent.parent.resolve() os.chdir(run_cwd) @@ -351,12 +351,12 @@ def test_sort(): try: with test_path.open("w") as f: f.write(code) - + code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/bubble_sort.py").resolve() tests_root = Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/" project_root_path = (Path(__file__).parent / "..").resolve() original_cwd = Path.cwd() run_cwd = Path(__file__).parent.parent.resolve() - func = FunctionToOptimize(function_name="sorter", parents=[], file_path=Path("module.py")) + func = FunctionToOptimize(function_name="sorter", parents=[], file_path=code_path) os.chdir(run_cwd) success, new_test = inject_profiling_into_existing_test( test_path, @@ -394,18 +394,14 @@ def test_sort(): # Overwrite old test with new instrumented 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_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="pytest", + pytest_cmd="pytest", ) - + func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config) test_env = os.environ.copy() test_env["CODEFLASH_TEST_ITERATION"] = "0" test_env["CODEFLASH_LOOP_INDEX"] = "1" @@ -420,7 +416,7 @@ def test_sort(): ) ] ) - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -455,7 +451,7 @@ def test_sort(): with test_path_perf.open("w") as f: f.write(new_perf_test) - test_results_perf, _ = opt.run_and_parse_tests( + test_results_perf, _ = func_optimizer.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, test_files=test_files, @@ -556,6 +552,7 @@ def test_sort_parametrized(input, expected_output): assert output == expected_output """ ) + code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/bubble_sort.py").resolve() test_path = ( Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_parametrized_results_temp.py" @@ -573,7 +570,7 @@ def test_sort_parametrized(input, expected_output): original_cwd = Path.cwd() run_cwd = Path(__file__).parent.parent.resolve() - func = FunctionToOptimize(function_name="sorter", parents=[], file_path=Path("module.py")) + func = FunctionToOptimize(function_name="sorter", parents=[], file_path=code_path) os.chdir(run_cwd) success, new_test = inject_profiling_into_existing_test( test_path, [CodePosition(14, 13)], func, project_root_path, "pytest", mode=TestingMode.BEHAVIOR @@ -614,19 +611,15 @@ def test_sort_parametrized(input, expected_output): ) ] ) - - 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_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="pytest", + pytest_cmd="pytest", ) - test_results, coverage_data = opt.run_and_parse_tests( + func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config) + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -668,7 +661,7 @@ def test_sort_parametrized(input, expected_output): assert test_results[2].runtime > 0 assert test_results[2].did_pass - test_results_perf, coverage_data = opt.run_and_parse_tests( + test_results_perf, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, test_files=test_files, @@ -783,7 +776,7 @@ def test_sort_parametrized_loop(input, expected_output): assert output == expected_output """ ) - + code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/bubble_sort.py").resolve() test_path = ( Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_parametrized_loop_results_temp.py" @@ -805,7 +798,7 @@ def test_sort_parametrized_loop(input, expected_output): original_cwd = Path.cwd() run_cwd = Path(__file__).parent.parent.resolve() - func = FunctionToOptimize(function_name="sorter", parents=[], file_path=Path("module.py")) + func = FunctionToOptimize(function_name="sorter", parents=[], file_path=code_path) os.chdir(run_cwd) success, new_test = inject_profiling_into_existing_test( test_path, [CodePosition(15, 17)], func, project_root_path, "pytest", mode=TestingMode.BEHAVIOR @@ -857,18 +850,16 @@ def test_sort_parametrized_loop(input, expected_output): ) ] ) - 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_results, coverage_data = opt.run_and_parse_tests( + + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="pytest", + pytest_cmd="pytest", + ) + func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config) + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -944,7 +935,7 @@ def test_sort_parametrized_loop(input, expected_output): assert test_results[5].runtime > 0 assert test_results[5].did_pass - test_results, _ = opt.run_and_parse_tests( + test_results, _ = func_optimizer.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, test_files=test_files, @@ -1097,7 +1088,7 @@ def test_sort(): assert output == expected_output """ ) - + code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/bubble_sort.py").resolve() test_path = ( Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_perfinjector_bubble_sort_loop_results_temp.py" @@ -1119,7 +1110,7 @@ def test_sort(): run_cwd = Path(__file__).parent.parent.resolve() original_cwd = Path.cwd() - func = FunctionToOptimize(function_name="sorter", parents=[], file_path=Path("module.py")) + func = FunctionToOptimize(function_name="sorter", parents=[], file_path=code_path) os.chdir(str(run_cwd)) success, new_test_behavior = inject_profiling_into_existing_test( test_path, [CodePosition(11, 17)], func, project_root_path, "pytest", mode=TestingMode.BEHAVIOR @@ -1169,18 +1160,15 @@ def test_sort(): ] ) - 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_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="pytest", + pytest_cmd="pytest", ) - test_results, coverage_data = opt.run_and_parse_tests( + func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config) + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -1222,7 +1210,7 @@ def test_sort(): ) assert test_results[2].runtime > 0 assert test_results[2].did_pass - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, test_files=test_files, @@ -1358,7 +1346,7 @@ def test_sort(self): self.assertEqual(output, list(range(50))) """ ) - + code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/bubble_sort.py").resolve() test_path = ( Path(__file__).parent.resolve() / "../code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_results_temp.py" @@ -1380,7 +1368,7 @@ def test_sort(self): run_cwd = Path(__file__).parent.parent.resolve() original_cwd = Path.cwd() - func = FunctionToOptimize(function_name="sorter", parents=[], file_path=Path("module.py")) + func = FunctionToOptimize(function_name="sorter", parents=[], file_path=code_path) os.chdir(run_cwd) success, new_test_behavior = inject_profiling_into_existing_test( test_path, @@ -1440,18 +1428,15 @@ def test_sort(self): ) ] ) - opt = Optimizer( - Namespace( - project_root=project_root_path, - disable_telemetry=True, - tests_root=tests_root, - test_framework="unittest", - pytest_cmd="pytest", - experiment_id=None, - test_project_root=project_root_path, - ) - ) - test_results, coverage_data = opt.run_and_parse_tests( + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="unittest", + pytest_cmd="pytest", + ) + func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config) + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -1493,7 +1478,7 @@ def test_sort(self): ) assert test_results[2].runtime > 0 assert test_results[2].did_pass - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, test_files=test_files, @@ -1617,7 +1602,7 @@ def test_sort(self, input, expected_output): self.assertEqual(output, expected_output) """ ) - + code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/bubble_sort.py").resolve() test_path = ( Path(__file__).parent.resolve() / "../code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_parametrized_results_temp.py" @@ -1638,7 +1623,7 @@ def test_sort(self, input, expected_output): run_cwd = Path(__file__).parent.parent.resolve() original_cwd = Path.cwd() - func = FunctionToOptimize(function_name="sorter", parents=[], file_path=Path("module.py")) + func = FunctionToOptimize(function_name="sorter", parents=[], file_path=code_path) os.chdir(run_cwd) success, new_test_behavior = inject_profiling_into_existing_test( test_path, [CodePosition(16, 17)], func, project_root_path, "unittest", mode=TestingMode.BEHAVIOR @@ -1690,18 +1675,15 @@ def test_sort(self, input, expected_output): ) ] ) - opt = Optimizer( - Namespace( - project_root=project_root_path, - disable_telemetry=True, - tests_root=tests_root, - test_framework="unittest", - pytest_cmd="pytest", - experiment_id=None, - test_project_root=project_root_path, - ) - ) - test_results, coverage_data = opt.run_and_parse_tests( + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="unittest", + pytest_cmd="pytest", + ) + func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config) + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -1744,7 +1726,7 @@ def test_sort(self, input, expected_output): assert test_results[2].runtime > 0 assert test_results[2].did_pass - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, test_files=test_files, @@ -1874,7 +1856,7 @@ def test_sort(self): self.assertEqual(output, expected_output) """ ) - + code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/bubble_sort.py").resolve() test_path = ( Path(__file__).parent.resolve() / "../code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_loop_results_temp.py" @@ -1896,7 +1878,7 @@ def test_sort(self): run_cwd = Path(__file__).parent.parent.resolve() original_cwd = Path.cwd() - func = FunctionToOptimize(function_name="sorter", parents=[], file_path=Path("module.py")) + func = FunctionToOptimize(function_name="sorter", parents=[], file_path=code_path) os.chdir(run_cwd) success, new_test_behavior = inject_profiling_into_existing_test( test_path, [CodePosition(14, 21)], func, project_root_path, "unittest", mode=TestingMode.BEHAVIOR @@ -1944,19 +1926,15 @@ def test_sort(self): ) ] ) - - opt = Optimizer( - Namespace( - project_root=project_root_path, - disable_telemetry=True, - tests_root=tests_root, - test_framework="unittest", - pytest_cmd="pytest", - experiment_id=None, - test_project_root=project_root_path, - ) + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="unittest", + pytest_cmd="pytest", ) - test_results, coverage_data = opt.run_and_parse_tests( + func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config) + test_results, coverage_data = func_optimizer.run_and_parse_tests( test_env=test_env, testing_type=TestingMode.BEHAVIOR, test_files=test_files, @@ -1999,7 +1977,7 @@ def test_sort(self): assert test_results[2].runtime > 0 assert test_results[2].did_pass - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( test_env=test_env, testing_type=TestingMode.PERFORMANCE, test_files=test_files, @@ -2127,7 +2105,7 @@ def test_sort(self, input, expected_output): self.assertEqual(output, expected_output) """ ) - + code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/bubble_sort.py").resolve() test_path = ( Path(__file__).parent.resolve() / "../code_to_optimize/tests/unittest/test_perfinjector_bubble_sort_unittest_parametrized_loop_results_temp.py" @@ -2148,7 +2126,7 @@ def test_sort(self, input, expected_output): run_cwd = Path(__file__).parent.parent.resolve() original_cwd = Path.cwd() - f = FunctionToOptimize(function_name="sorter", file_path=Path("module.py"), parents=[]) + f = FunctionToOptimize(function_name="sorter", file_path=code_path, parents=[]) os.chdir(run_cwd) success, new_test_behavior = inject_profiling_into_existing_test( test_path, [CodePosition(17, 21)], f, project_root_path, "unittest", mode=TestingMode.BEHAVIOR @@ -2197,18 +2175,15 @@ def test_sort(self, input, expected_output): ) ] ) - opt = Optimizer( - Namespace( - project_root=project_root_path, - disable_telemetry=True, - tests_root=tests_root, - test_framework="unittest", - pytest_cmd="pytest", - experiment_id=None, - test_project_root=project_root_path, - ) - ) - test_results, coverage_data = opt.run_and_parse_tests( + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="unittest", + pytest_cmd="pytest", + ) + func_optimizer = FunctionOptimizer(function_to_optimize=f, test_cfg=test_config) + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -2283,7 +2258,7 @@ def test_sort(self, input, expected_output): ) assert test_results[5].runtime > 0 assert test_results[5].did_pass - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, test_files=test_files, @@ -2776,7 +2751,7 @@ def test_code_replacement10() -> None: func = FunctionToOptimize( function_name="get_code_optimization_context", parents=[FunctionParent("Optimizer", "ClassDef")], - file_path=Path("module.py"), + file_path=Path(f.name), ) original_cwd = Path.cwd() run_cwd = Path(__file__).parent.parent.resolve() @@ -2824,7 +2799,7 @@ def test_sleepfunc_sequence_short(n, expected_total_sleep_time): assert output == expected_total_sleep_time """ ) - + code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/sleeptime.py").resolve() test_path = ( Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_time_correction_instrumentation_temp.py" @@ -2837,7 +2812,7 @@ def test_sleepfunc_sequence_short(n, expected_total_sleep_time): project_root_path = (Path(__file__).parent.resolve() / "../").resolve() original_cwd = Path.cwd() run_cwd = Path(__file__).parent.parent.resolve() - func = FunctionToOptimize(function_name="accurate_sleepfunc", parents=[], file_path=Path("module.py")) + func = FunctionToOptimize(function_name="accurate_sleepfunc", parents=[], file_path=code_path) os.chdir(run_cwd) success, new_test = inject_profiling_into_existing_test( test_path, [CodePosition(8, 13)], func, project_root_path, "pytest", mode=TestingMode.PERFORMANCE @@ -2858,17 +2833,14 @@ def test_sleepfunc_sequence_short(n, expected_total_sleep_time): 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_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="pytest", + pytest_cmd="pytest", ) + func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config) test_files = TestFiles( test_files=[ TestFile( @@ -2879,7 +2851,7 @@ def test_sleepfunc_sequence_short(n, expected_total_sleep_time): ) ] ) - test_results, coverage_data = opt.run_and_parse_tests( + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, test_files=test_files, @@ -2946,7 +2918,7 @@ def test_sleepfunc_sequence_short(self, n, expected_total_sleep_time): output = codeflash_wrap(accurate_sleepfunc, '{module_path}', 'TestPigLatin', 'test_sleepfunc_sequence_short', 'accurate_sleepfunc', '0', codeflash_loop_index, n) """ ) - + code_path = (Path(__file__).parent.resolve() / "../code_to_optimize/sleeptime.py").resolve() test_path = ( Path(__file__).parent.resolve() / "../code_to_optimize/tests/unittest/test_time_correction_instrumentation_unittest_temp.py" @@ -2959,7 +2931,7 @@ def test_sleepfunc_sequence_short(self, n, expected_total_sleep_time): project_root_path = (Path(__file__).parent.resolve() / "../").resolve() original_cwd = Path.cwd() run_cwd = Path(__file__).parent.parent.resolve() - func = FunctionToOptimize(function_name="accurate_sleepfunc", parents=[], file_path=Path("module.py")) + func = FunctionToOptimize(function_name="accurate_sleepfunc", parents=[], file_path=code_path) os.chdir(run_cwd) success, new_test = inject_profiling_into_existing_test( test_path, [CodePosition(12, 17)], func, project_root_path, "unittest", mode=TestingMode.PERFORMANCE @@ -2980,17 +2952,6 @@ def test_sleepfunc_sequence_short(self, n, expected_total_sleep_time): 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="unittest", - pytest_cmd="pytest", - experiment_id=None, - test_project_root=project_root_path, - ) - ) test_files = TestFiles( test_files=[ TestFile( @@ -3009,7 +2970,15 @@ def test_sleepfunc_sequence_short(self, n, expected_total_sleep_time): ) ] ) - test_results, coverage_data = opt.run_and_parse_tests( + test_config = TestConfig( + tests_root=tests_root, + tests_project_rootdir=project_root_path, + project_root_path=project_root_path, + test_framework="unittest", + pytest_cmd="pytest", + ) + func_optimizer = FunctionOptimizer(function_to_optimize=func, test_cfg=test_config) + test_results, coverage_data = func_optimizer.run_and_parse_tests( testing_type=TestingMode.PERFORMANCE, test_env=test_env, test_files=test_files, diff --git a/tests/test_instrumentation_run_results_aiservice.py b/tests/test_instrumentation_run_results_aiservice.py index c239f21bb..ee237cfca 100644 --- a/tests/test_instrumentation_run_results_aiservice.py +++ b/tests/test_instrumentation_run_results_aiservice.py @@ -118,9 +118,14 @@ def test_single_element_list(): ) # Init paths - tests_root = (Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/").resolve() - test_path = tests_root / "test_aiservice_behavior_results_temp.py" - test_path_perf = tests_root / "test_aiservice_behavior_results_perf_temp.py" + test_path = ( + Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/test_aiservice_behavior_results_temp.py" + ).resolve() + test_path_perf = ( + Path(__file__).parent.resolve() + / "../code_to_optimize/tests/pytest/test_aiservice_behavior_results_perf_temp.py" + ).resolve() + tests_root = Path(__file__).parent.resolve() / "../code_to_optimize/tests/pytest/" project_root_path = (Path(__file__).parent / "..").resolve() run_cwd = Path(__file__).parent.parent.resolve() os.chdir(run_cwd) @@ -161,7 +166,9 @@ def test_single_element_list(): ] ) a = BubbleSorter() - test_results, coverage_data = opt.run_and_parse_tests( + function_to_optimize = FunctionToOptimize("sorter", fto_path, [FunctionParent("BubbleSorter", "ClassDef")]) + func_opt = opt.create_function_optimizer(function_to_optimize) + test_results, coverage_data = func_opt.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -194,7 +201,8 @@ def sorter(self, arr): return arr """ fto_path.write_text(optimized_code_mutated_attr, "utf-8") - test_results_mutated_attr, coverage_data = opt.run_and_parse_tests( + func_opt = opt.create_function_optimizer(function_to_optimize) + test_results_mutated_attr, coverage_data = func_opt.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -293,7 +301,8 @@ def test_single_element_list(): ) ] ) - test_results, coverage_data = opt.run_and_parse_tests( + func_opt = opt.create_function_optimizer(function_to_optimize) + test_results, coverage_data = func_opt.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -357,7 +366,8 @@ def sorter(self, arr): test_project_root=project_root_path, ) ) - test_results_mutated_attr, coverage_data = opt.run_and_parse_tests( + func_opt = opt.create_function_optimizer(function_to_optimize) + test_results_mutated_attr, coverage_data = func_opt.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files, @@ -403,7 +413,8 @@ def sorter(self, arr): test_project_root=project_root_path, ) ) - test_results_new_attr, coverage_data = opt.run_and_parse_tests( + func_opt = opt.create_function_optimizer(function_to_optimize) + test_results_new_attr, coverage_data = func_opt.run_and_parse_tests( testing_type=TestingMode.BEHAVIOR, test_env=test_env, test_files=test_files,