Split test generation and optimizations generation into separate methods (#957)

Author: KRRT7

codeflash/api/aiservice.py

@@ -153,7 +153,6 @@ def optimize_python_code(  # noqa: D417
 
         if response.status_code == 200:
             optimizations_json = response.json()["optimizations"]
-            logger.info(f"!lsp|Generated {len(optimizations_json)} candidate optimizations.")
             console.rule()
             end_time = time.perf_counter()
             logger.debug(f"!lsp|Generating possible optimizations took {end_time - start_time:.2f} seconds.")

codeflash/lsp/features/perform_optimization.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+import concurrent.futures
 import contextlib
+import contextvars
 import os
 from typing import TYPE_CHECKING
 
-from codeflash.cli_cmds.console import code_print
+from codeflash.cli_cmds.console import code_print, logger
 from codeflash.code_utils.git_worktree_utils import create_diff_patch_from_worktree
 from codeflash.either import is_successful
 
@@ -44,24 +46,48 @@ def sync_perform_optimization(server: CodeflashLanguageServer, cancel_event: thr
         function_optimizer.function_to_tests = function_to_tests
 
     abort_if_cancelled(cancel_event)
-    test_setup_result = function_optimizer.generate_and_instrument_tests(
-        code_context, should_run_experiment=should_run_experiment
-    )
+
+    ctx_tests = contextvars.copy_context()
+    ctx_opts = contextvars.copy_context()
+
+    def run_generate_tests():  # noqa: ANN202
+        return function_optimizer.generate_and_instrument_tests(code_context)
+
+    def run_generate_optimizations():  # noqa: ANN202
+        return function_optimizer.generate_optimizations(
+            read_writable_code=code_context.read_writable_code,
+            read_only_context_code=code_context.read_only_context_code,
+            run_experiment=should_run_experiment,
+        )
+
+    future_tests = function_optimizer.executor.submit(ctx_tests.run, run_generate_tests)
+    future_optimizations = function_optimizer.executor.submit(ctx_opts.run, run_generate_optimizations)
+
+    logger.info(f"loading|Generating new tests and optimizations for function '{params.functionName}'...")
+    concurrent.futures.wait([future_tests, future_optimizations])
+
+    test_setup_result = future_tests.result()
+    optimization_result = future_optimizations.result()
+
     abort_if_cancelled(cancel_event)
     if not is_successful(test_setup_result):
         return {"functionName": params.functionName, "status": "error", "message": test_setup_result.failure()}
+    if not is_successful(optimization_result):
+        return {"functionName": params.functionName, "status": "error", "message": optimization_result.failure()}
+
     (
         generated_tests,
         function_to_concolic_tests,
         concolic_test_str,
-        optimizations_set,
         generated_test_paths,
         generated_perf_test_paths,
         instrumented_unittests_created_for_function,
         original_conftest_content,
-        function_references,
     ) = test_setup_result.unwrap()
 
+    optimizations_set, function_references = optimization_result.unwrap()
+
+    logger.info(f"Generated '{len(optimizations_set.control)}' candidate optimizations.")
     baseline_setup_result = function_optimizer.setup_and_establish_baseline(
         code_context=code_context,
         original_helper_code=original_helper_code,

codeflash/optimization/function_optimizer.py

@@ -238,9 +238,6 @@ def __init__(
         self.function_benchmark_timings = function_benchmark_timings if function_benchmark_timings else {}
         self.total_benchmark_timings = total_benchmark_timings if total_benchmark_timings else {}
         self.replay_tests_dir = replay_tests_dir if replay_tests_dir else None
-        self.generate_and_instrument_tests_results: (
-            tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet] | None
-        ) = None
         n_tests = N_TESTS_TO_GENERATE_EFFECTIVE
         self.executor = concurrent.futures.ThreadPoolExecutor(
             max_workers=n_tests + 3 if self.experiment_id is None else n_tests + 4
@@ -275,21 +272,20 @@ def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[P
         return Success((should_run_experiment, code_context, original_helper_code))
 
     def generate_and_instrument_tests(
-        self, code_context: CodeOptimizationContext, *, should_run_experiment: bool
+        self, code_context: CodeOptimizationContext
     ) -> Result[
         tuple[
             GeneratedTestsList,
             dict[str, set[FunctionCalledInTest]],
             str,
-            OptimizationSet,
             list[Path],
             list[Path],
             set[Path],
             dict | None,
-            str,
-        ]
+        ],
+        str,
     ]:
-        """Generate and instrument tests, returning all necessary data for optimization."""
+        """Generate and instrument tests for the function."""
         n_tests = N_TESTS_TO_GENERATE_EFFECTIVE
         generated_test_paths = [
             get_test_file_path(
@@ -304,34 +300,17 @@ def generate_and_instrument_tests(
             for test_index in range(n_tests)
         ]
 
-        with progress_bar(
-            f"Generating new tests and optimizations for function '{self.function_to_optimize.function_name}'",
-            transient=True,
-            revert_to_print=bool(get_pr_number()),
-        ):
-            generated_results = self.generate_tests_and_optimizations(
-                testgen_context=code_context.testgen_context,
-                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,
-            )
+        test_results = self.generate_tests(
+            testgen_context=code_context.testgen_context,
+            helper_functions=code_context.helper_functions,
+            generated_test_paths=generated_test_paths,
+            generated_perf_test_paths=generated_perf_test_paths,
+        )
 
-        if not is_successful(generated_results):
-            return Failure(generated_results.failure())
+        if not is_successful(test_results):
+            return Failure(test_results.failure())
 
-        generated_tests: GeneratedTestsList
-        optimizations_set: OptimizationSet
-        (
-            count_tests,
-            generated_tests,
-            function_to_concolic_tests,
-            concolic_test_str,
-            optimizations_set,
-            function_references,
-        ) = generated_results.unwrap()
+        count_tests, generated_tests, function_to_concolic_tests, concolic_test_str = test_results.unwrap()
 
         for i, generated_test in enumerate(generated_tests.generated_tests):
             with generated_test.behavior_file_path.open("w", encoding="utf8") as f:
@@ -372,12 +351,10 @@ def generate_and_instrument_tests(
                 generated_tests,
                 function_to_concolic_tests,
                 concolic_test_str,
-                optimizations_set,
                 generated_test_paths,
                 generated_perf_test_paths,
                 instrumented_unittests_created_for_function,
                 original_conftest_content,
-                function_references,
             )
         )
 
@@ -395,24 +372,45 @@ def optimize_function(self) -> Result[BestOptimization, str]:
             function_name=self.function_to_optimize.function_name,
         )
 
-        test_setup_result = self.generate_and_instrument_tests(  # also generates optimizations
-            code_context, should_run_experiment=should_run_experiment
-        )
+        with progress_bar(
+            f"Generating new tests and optimizations for function '{self.function_to_optimize.function_name}'",
+            transient=True,
+            revert_to_print=bool(get_pr_number()),
+        ):
+            console.rule()
+            # Generate tests and optimizations in parallel
+            future_tests = self.executor.submit(self.generate_and_instrument_tests, code_context)
+            future_optimizations = self.executor.submit(
+                self.generate_optimizations,
+                read_writable_code=code_context.read_writable_code,
+                read_only_context_code=code_context.read_only_context_code,
+                run_experiment=should_run_experiment,
+            )
+
+            concurrent.futures.wait([future_tests, future_optimizations])
+
+            test_setup_result = future_tests.result()
+            optimization_result = future_optimizations.result()
+            console.rule()
+
         if not is_successful(test_setup_result):
             return Failure(test_setup_result.failure())
 
+        if not is_successful(optimization_result):
+            return Failure(optimization_result.failure())
+
         (
             generated_tests,
             function_to_concolic_tests,
             concolic_test_str,
-            optimizations_set,
             generated_test_paths,
             generated_perf_test_paths,
             instrumented_unittests_created_for_function,
             original_conftest_content,
-            function_references,
         ) = test_setup_result.unwrap()
 
+        optimizations_set, function_references = optimization_result.unwrap()
+
         baseline_setup_result = self.setup_and_establish_baseline(
             code_context=code_context,
             original_helper_code=original_helper_code,
@@ -1109,28 +1107,78 @@ def instrument_existing_tests(self, function_to_all_tests: dict[str, set[Functio
             console.rule()
         return unique_instrumented_test_files
 
-    def generate_tests_and_optimizations(
+    def generate_tests(
         self,
         testgen_context: CodeStringsMarkdown,
-        read_writable_code: CodeStringsMarkdown,
-        read_only_context_code: str,
         helper_functions: list[FunctionSource],
         generated_test_paths: list[Path],
         generated_perf_test_paths: list[Path],
-        run_experiment: bool = False,  # noqa: FBT001, FBT002
-    ) -> Result[tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet], str, str]:
+    ) -> Result[tuple[int, GeneratedTestsList, dict[str, set[FunctionCalledInTest]], str], str]:
+        """Generate unit tests and concolic tests for the function."""
         n_tests = N_TESTS_TO_GENERATE_EFFECTIVE
         assert len(generated_test_paths) == n_tests
-        console.rule()
-        # Submit the test generation task as future
+
+        # Submit test generation tasks
         future_tests = self.submit_test_generation_tasks(
             self.executor,
             testgen_context.markdown,
             [definition.fully_qualified_name for definition in helper_functions],
             generated_test_paths,
             generated_perf_test_paths,
         )
+
+        future_concolic_tests = self.executor.submit(
+            generate_concolic_tests, self.test_cfg, self.args, self.function_to_optimize, self.function_to_optimize_ast
+        )
+
+        # Wait for test futures to complete
+        concurrent.futures.wait([*future_tests, future_concolic_tests])
+
+        # 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()
+        count_tests = len(tests)
+        if concolic_test_str:
+            count_tests += 1
+
+        logger.info(f"!lsp|Generated '{count_tests}' tests for '{self.function_to_optimize.function_name}'")
+
+        generated_tests = GeneratedTestsList(generated_tests=tests)
+        return Success((count_tests, generated_tests, function_to_concolic_tests, concolic_test_str))
+
+    def generate_optimizations(
+        self,
+        read_writable_code: CodeStringsMarkdown,
+        read_only_context_code: str,
+        run_experiment: bool = False,  # noqa: FBT001, FBT002
+    ) -> Result[tuple[OptimizationSet, str], str]:
+        """Generate optimization candidates for the function."""
         n_candidates = N_CANDIDATES_EFFECTIVE
+
         future_optimization_candidates = self.executor.submit(
             self.aiservice_client.optimize_python_code,
             read_writable_code.markdown,
@@ -1140,11 +1188,7 @@ def generate_tests_and_optimizations(
             ExperimentMetadata(id=self.experiment_id, group="control") if run_experiment else None,
             is_async=self.function_to_optimize.is_async,
         )
-        future_candidates_exp = None
 
-        future_concolic_tests = self.executor.submit(
-            generate_concolic_tests, self.test_cfg, self.args, self.function_to_optimize, self.function_to_optimize_ast
-        )
         future_references = self.executor.submit(
             get_opt_review_metrics,
             self.function_to_optimize_source_code,
@@ -1153,7 +1197,10 @@ def generate_tests_and_optimizations(
             self.project_root,
             self.test_cfg.tests_root,
         )
-        futures = [*future_tests, future_optimization_candidates, future_concolic_tests, future_references]
+
+        futures = [future_optimization_candidates, future_references]
+        future_candidates_exp = None
+
         if run_experiment:
             future_candidates_exp = self.executor.submit(
                 self.local_aiservice_client.optimize_python_code,
@@ -1166,63 +1213,20 @@ def generate_tests_and_optimizations(
             )
             futures.append(future_candidates_exp)
 
-        # Wait for all futures to complete
+        # Wait for optimization futures to complete
         concurrent.futures.wait(futures)
 
         # Retrieve results
         candidates: list[OptimizedCandidate] = future_optimization_candidates.result()
-        logger.info(f"lsp|Generated '{len(candidates)}' candidate optimizations.")
-        console.rule()
+        logger.info(f"!lsp|Generated '{len(candidates)}' candidate optimizations.")
 
         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()
         function_references = future_references.result()
-        count_tests = len(tests)
-        if concolic_test_str:
-            count_tests += 1
 
-        logger.info(f"Generated '{count_tests}' tests for '{self.function_to_optimize.function_name}'")
-        console.rule()
-        generated_tests = GeneratedTestsList(generated_tests=tests)
-        result = (
-            count_tests,
-            generated_tests,
-            function_to_concolic_tests,
-            concolic_test_str,
-            OptimizationSet(control=candidates, experiment=candidates_experiment),
-            function_references,
-        )
-        self.generate_and_instrument_tests_results = result
-        return Success(result)
+        return Success((OptimizationSet(control=candidates, experiment=candidates_experiment), function_references))
 
     def setup_and_establish_baseline(
         self,