Merge pull request #1461 from codeflash-ai/aseembits93-patch-1

Restore concurrency in testgen and candidate generation

Author: KRRT7

codeflash/optimization/function_optimizer.py

@@ -13,7 +13,6 @@
 from typing import TYPE_CHECKING, Callable
 
 import libcst as cst
-import sentry_sdk
 from rich.console import Group
 from rich.panel import Panel
 from rich.syntax import Syntax
@@ -680,30 +679,6 @@ def optimize_function(self) -> Result[BestOptimization, str]:
         ):
             console.rule()
             new_code_context = code_context
-            if (
-                self.is_numerical_code and not self.args.no_jit_opts
-            ):  # if the code is numerical in nature (uses numpy/tensorflow/math/pytorch/jax)
-                jit_compiled_opt_candidate = self.aiservice_client.get_jit_rewritten_code(
-                    code_context.read_writable_code.markdown, self.function_trace_id
-                )
-                if jit_compiled_opt_candidate:  # jit rewrite was successful
-                    # write files
-                    # Try to replace function with optimized code
-                    self.replace_function_and_helpers_with_optimized_code(
-                        code_context=code_context,
-                        optimized_code=jit_compiled_opt_candidate[0].source_code,
-                        original_helper_code=original_helper_code,
-                    )
-                    # get code context
-                    try:
-                        new_code_context = self.get_code_optimization_context().unwrap()
-                    except Exception as e:
-                        sentry_sdk.capture_exception(e)
-                        logger.debug("!lsp|Getting new code context failed, revert to original one")
-                    # unwrite files
-                    self.write_code_and_helpers(
-                        self.function_to_optimize_source_code, original_helper_code, self.function_to_optimize.file_path
-                    )
             # Generate tests and optimizations in parallel
             future_tests = self.executor.submit(self.generate_and_instrument_tests, new_code_context)
             future_optimizations = self.executor.submit(