don't reapply too early

Author: KRRT7

codeflash/code_utils/instrument_existing_tests.py

@@ -77,10 +77,10 @@ def find_and_update_line_node(
                 call_node = node
                 if isinstance(node.func, ast.Name):
                     function_name = node.func.id
-                    
+
                     if self.function_object.is_async:
                         return [test_node]
-                    
+
                     node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
                     node.args = [
                         ast.Name(id=function_name, ctx=ast.Load()),
@@ -104,7 +104,7 @@ def find_and_update_line_node(
                     if function_to_test == self.function_object.function_name:
                         if self.function_object.is_async:
                             return [test_node]
-                        
+
                         function_name = ast.unparse(node.func)
                         node.func = ast.Name(id="codeflash_wrap", ctx=ast.Load())
                         node.args = [
@@ -359,11 +359,10 @@ def inject_profiling_into_existing_test(
     tests_project_root: Path,
     test_framework: str,
     mode: TestingMode = TestingMode.BEHAVIOR,
-    source_module_path: Path | None = None,
 ) -> tuple[bool, str | None]:
     if function_to_optimize.is_async:
         return False, None
-        
+
     with test_path.open(encoding="utf8") as f:
         test_code = f.read()
     try:
@@ -390,7 +389,7 @@ def inject_profiling_into_existing_test(
     if test_framework == "unittest":
         new_imports.append(ast.Import(names=[ast.alias(name="timeout_decorator")]))
     additional_functions = [create_wrapper_function(mode)]
-    
+
     tree.body = [*new_imports, *additional_functions, *tree.body]
     return True, isort.code(ast.unparse(tree), float_to_top=True)
 
@@ -925,5 +924,3 @@ def add_async_decorator_to_function(
 def create_instrumented_source_module_path(source_path: Path, temp_dir: Path) -> Path:
     instrumented_filename = f"instrumented_{source_path.name}"
     return temp_dir / instrumented_filename
-
-

codeflash/optimization/function_optimizer.py

@@ -893,15 +893,14 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, set[Functio
                 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,
-                    source_module_path=self.function_to_optimize.file_path,
                 )
                 if not success:
                     continue
@@ -912,7 +911,6 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, set[Functio
                     function_to_optimize=self.function_to_optimize,
                     tests_project_root=self.test_cfg.tests_project_rootdir,
                     test_framework=self.args.test_framework,
-                    source_module_path=self.function_to_optimize.file_path,
                 )
                 if not success:
                     continue
@@ -1368,15 +1366,18 @@ def establish_original_code_baseline(
             test_env = self.get_test_env(codeflash_loop_index=0, codeflash_test_iteration=0, codeflash_tracer_disable=1)
 
             if self.function_to_optimize.is_async:
-                from codeflash.code_utils.instrument_existing_tests import instrument_source_module_with_async_decorators
+                from codeflash.code_utils.instrument_existing_tests import (
+                    instrument_source_module_with_async_decorators,
+                )
+
                 success, instrumented_source = instrument_source_module_with_async_decorators(
                     self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.BEHAVIOR
                 )
                 if success and instrumented_source:
                     with self.function_to_optimize.file_path.open("w", encoding="utf8") as f:
                         f.write(instrumented_source)
                     logger.debug(f"Applied async instrumentation to {self.function_to_optimize.file_path}")
-            
+
             # Instrument codeflash capture
             try:
                 instrument_codeflash_capture(
@@ -1392,9 +1393,10 @@ def establish_original_code_baseline(
                     code_context=code_context,
                 )
             finally:
-                self.write_code_and_helpers(
-                    self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
-                )
+                if not self.function_to_optimize.is_async:
+                    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."
@@ -1408,18 +1410,22 @@ def establish_original_code_baseline(
                     code_context=code_context, original_helper_code=original_helper_code, candidate_index=0
                 )
                 console.rule()
-                
-                # For async functions, instrument at definition site for performance benchmarking
+
                 if self.function_to_optimize.is_async:
-                    from codeflash.code_utils.instrument_existing_tests import instrument_source_module_with_async_decorators
+                    from codeflash.code_utils.instrument_existing_tests import (
+                        instrument_source_module_with_async_decorators,
+                    )
+
                     success, instrumented_source = instrument_source_module_with_async_decorators(
                         self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.PERFORMANCE
                     )
                     if success and instrumented_source:
                         with self.function_to_optimize.file_path.open("w", encoding="utf8") as f:
                             f.write(instrumented_source)
-                        logger.debug(f"Applied async performance instrumentation to {self.function_to_optimize.file_path}")
-                
+                        logger.debug(
+                            f"Applied async performance instrumentation to {self.function_to_optimize.file_path}"
+                        )
+
                 try:
                     benchmarking_results, _ = self.run_and_parse_tests(
                         testing_type=TestingMode.PERFORMANCE,
@@ -1431,10 +1437,11 @@ def establish_original_code_baseline(
                         code_context=code_context,
                     )
                 finally:
-                    # Restore original source if we instrumented it
                     if self.function_to_optimize.is_async:
                         self.write_code_and_helpers(
-                            self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
+                            self.function_to_optimize_source_code,
+                            original_helper_code,
+                            self.function_to_optimize.file_path,
                         )
             else:
                 benchmarking_results = TestResults()
@@ -1570,15 +1577,20 @@ def run_optimized_candidate(
             if test_framework == "pytest":
                 # For async functions, instrument at definition site for performance benchmarking
                 if self.function_to_optimize.is_async:
-                    from codeflash.code_utils.instrument_existing_tests import instrument_source_module_with_async_decorators
+                    from codeflash.code_utils.instrument_existing_tests import (
+                        instrument_source_module_with_async_decorators,
+                    )
+
                     success, instrumented_source = instrument_source_module_with_async_decorators(
                         self.function_to_optimize.file_path, self.function_to_optimize, TestingMode.PERFORMANCE
                     )
                     if success and instrumented_source:
                         with self.function_to_optimize.file_path.open("w", encoding="utf8") as f:
                             f.write(instrumented_source)
-                        logger.debug(f"Applied async performance instrumentation to {self.function_to_optimize.file_path} for candidate {optimization_candidate_index}")
-                
+                        logger.debug(
+                            f"Applied async performance instrumentation to {self.function_to_optimize.file_path} for candidate {optimization_candidate_index}"
+                        )
+
                 try:
                     candidate_benchmarking_results, _ = self.run_and_parse_tests(
                         testing_type=TestingMode.PERFORMANCE,