LSP reduce no of candidates

Author: Saga4

codeflash/code_utils/config_consts.py

@@ -11,3 +11,26 @@
 MIN_TESTCASE_PASSED_THRESHOLD = 6
 REPEAT_OPTIMIZATION_PROBABILITY = 0.1
 DEFAULT_IMPORTANCE_THRESHOLD = 0.001
+
+# LSP-specific
+N_CANDIDATES_LSP = 3
+N_TESTS_TO_GENERATE_LSP = 1
+TOTAL_LOOPING_TIME_LSP = 5.0
+
+
+def get_n_candidates() -> int:
+    from codeflash.lsp.helpers import is_LSP_enabled
+
+    return N_CANDIDATES_LSP if is_LSP_enabled() else N_CANDIDATES
+
+
+def get_n_tests_to_generate() -> int:
+    from codeflash.lsp.helpers import is_LSP_enabled
+
+    return N_TESTS_TO_GENERATE_LSP if is_LSP_enabled() else N_TESTS_TO_GENERATE
+
+
+def get_total_looping_time() -> float:
+    from codeflash.lsp.helpers import is_LSP_enabled
+
+    return TOTAL_LOOPING_TIME_LSP if is_LSP_enabled() else TOTAL_LOOPING_TIME

codeflash/code_utils/git_utils.py

@@ -16,7 +16,7 @@
 from unidiff import PatchSet
 
 from codeflash.cli_cmds.console import logger
-from codeflash.code_utils.config_consts import N_CANDIDATES
+from codeflash.code_utils.config_consts import get_n_candidates
 
 if TYPE_CHECKING:
     from git import Repo
@@ -164,7 +164,7 @@ def create_git_worktrees(
 ) -> tuple[Path | None, list[Path]]:
     if git_root and worktree_root_dir:
         worktree_root = Path(tempfile.mkdtemp(dir=worktree_root_dir))
-        worktrees = [Path(tempfile.mkdtemp(dir=worktree_root)) for _ in range(N_CANDIDATES + 1)]
+        worktrees = [Path(tempfile.mkdtemp(dir=worktree_root)) for _ in range(get_n_candidates() + 1)]
         for worktree in worktrees:
             subprocess.run(["git", "worktree", "add", "-d", worktree], cwd=module_root, check=True)
     else:

codeflash/optimization/function_optimizer.py

@@ -42,10 +42,10 @@
 )
 from codeflash.code_utils.config_consts import (
     INDIVIDUAL_TESTCASE_TIMEOUT,
-    N_CANDIDATES,
-    N_TESTS_TO_GENERATE,
     REPEAT_OPTIMIZATION_PROBABILITY,
-    TOTAL_LOOPING_TIME,
+    get_n_candidates,
+    get_n_tests_to_generate,
+    get_total_looping_time,
 )
 from codeflash.code_utils.edit_generated_tests import (
     add_runtime_comments_to_generated_tests,
@@ -227,8 +227,9 @@ def __init__(
         self.generate_and_instrument_tests_results: (
             tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet] | None
         ) = None
+        n_tests = get_n_tests_to_generate()
         self.executor = concurrent.futures.ThreadPoolExecutor(
-            max_workers=N_TESTS_TO_GENERATE + 2 if self.experiment_id is None else N_TESTS_TO_GENERATE + 3
+            max_workers=n_tests + 2 if self.experiment_id is None else n_tests + 3
         )
 
     def can_be_optimized(self) -> Result[tuple[bool, CodeOptimizationContext, dict[Path, str]], str]:
@@ -278,17 +279,18 @@ def generate_and_instrument_tests(
         ]
     ]:
         """Generate and instrument tests, returning all necessary data for optimization."""
+        n_tests = get_n_tests_to_generate()
         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)
+            for test_index in range(n_tests)
         ]
         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)
+            for test_index in range(n_tests)
         ]
 
         with progress_bar(
@@ -971,7 +973,8 @@ def generate_tests_and_optimizations(
         generated_perf_test_paths: list[Path],
         run_experiment: bool = False,  # noqa: FBT001, FBT002
     ) -> Result[tuple[GeneratedTestsList, dict[str, set[FunctionCalledInTest]], OptimizationSet], str]:
-        assert len(generated_test_paths) == N_TESTS_TO_GENERATE
+        n_tests = get_n_tests_to_generate()
+        assert len(generated_test_paths) == n_tests
         console.rule()
         # Submit the test generation task as future
         future_tests = self.submit_test_generation_tasks(
@@ -981,12 +984,13 @@ def generate_tests_and_optimizations(
             generated_test_paths,
             generated_perf_test_paths,
         )
+        n_candidates = get_n_candidates()
         future_optimization_candidates = self.executor.submit(
             self.aiservice_client.optimize_python_code,
             read_writable_code.markdown,
             read_only_context_code,
             self.function_trace_id[:-4] + "EXP0" if run_experiment else self.function_trace_id,
-            N_CANDIDATES,
+            n_candidates,
             ExperimentMetadata(id=self.experiment_id, group="control") if run_experiment else None,
         )
         future_candidates_exp = None
@@ -1001,7 +1005,7 @@ def generate_tests_and_optimizations(
                 read_writable_code.markdown,
                 read_only_context_code,
                 self.function_trace_id[:-4] + "EXP1",
-                N_CANDIDATES,
+                n_candidates,
                 ExperimentMetadata(id=self.experiment_id, group="experiment"),
             )
             futures.append(future_candidates_exp)
@@ -1375,12 +1379,13 @@ def establish_original_code_baseline(
                 instrument_codeflash_capture(
                     self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root
                 )
+                total_looping_time = get_total_looping_time()
                 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,
+                    testing_time=total_looping_time,
                     enable_coverage=test_framework == "pytest",
                     code_context=code_context,
                 )
@@ -1407,15 +1412,15 @@ def establish_original_code_baseline(
                     test_env=test_env,
                     test_files=self.test_files,
                     optimization_iteration=0,
-                    testing_time=TOTAL_LOOPING_TIME,
+                    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:
+                    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)
@@ -1424,7 +1429,7 @@ def establish_original_code_baseline(
                         test_env=test_env,
                         test_files=self.test_files,
                         optimization_iteration=0,
-                        testing_time=TOTAL_LOOPING_TIME,
+                        testing_time=total_looping_time,
                         enable_coverage=False,
                         code_context=code_context,
                         unittest_loop_index=i + 1,
@@ -1514,12 +1519,13 @@ def run_optimized_candidate(
                     self.function_to_optimize, file_path_to_helper_classes, self.test_cfg.tests_root
                 )
 
+                total_looping_time = get_total_looping_time()
                 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,
+                    testing_time=total_looping_time,
                     enable_coverage=False,
                 )
             # Remove instrumentation
@@ -1548,7 +1554,7 @@ def run_optimized_candidate(
                     test_env=test_env,
                     test_files=self.test_files,
                     optimization_iteration=optimization_candidate_index,
-                    testing_time=TOTAL_LOOPING_TIME,
+                    testing_time=total_looping_time,
                     enable_coverage=False,
                 )
                 loop_count = (
@@ -1566,7 +1572,7 @@ def run_optimized_candidate(
                 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:
+                    if i >= 5 and time.time() - start_time >= get_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)
@@ -1575,7 +1581,7 @@ def run_optimized_candidate(
                         test_env=test_env,
                         test_files=self.test_files,
                         optimization_iteration=optimization_candidate_index,
-                        testing_time=TOTAL_LOOPING_TIME,
+                        testing_time=get_total_looping_time(),
                         unittest_loop_index=i + 1,
                     )
                     loop_count = i + 1
@@ -1614,7 +1620,7 @@ def run_and_parse_tests(
         test_env: dict[str, str],
         test_files: TestFiles,
         optimization_iteration: int,
-        testing_time: float = TOTAL_LOOPING_TIME,
+        testing_time: float = get_total_looping_time(),
         *,
         enable_coverage: bool = False,
         pytest_min_loops: int = 5,
@@ -1753,6 +1759,9 @@ def line_profiler_step(
         self, code_context: CodeOptimizationContext, original_helper_code: dict[Path, str], candidate_index: int
     ) -> dict:
         try:
+            logger.info("Running line profiling to identify performance bottlenecks…")
+            console.rule()
+
             test_env = self.get_test_env(
                 codeflash_loop_index=0, codeflash_test_iteration=candidate_index, codeflash_tracer_disable=1
             )
@@ -1762,7 +1771,7 @@ def line_profiler_step(
                 test_env=test_env,
                 test_files=self.test_files,
                 optimization_iteration=0,
-                testing_time=TOTAL_LOOPING_TIME,
+                testing_time=get_total_looping_time(),
                 enable_coverage=False,
                 code_context=code_context,
                 line_profiler_output_file=line_profiler_output_file,

codeflash/verification/test_runner.py

@@ -8,7 +8,7 @@
 from codeflash.cli_cmds.console import logger
 from codeflash.code_utils.code_utils import custom_addopts, get_run_tmp_file
 from codeflash.code_utils.compat import IS_POSIX, SAFE_SYS_EXECUTABLE
-from codeflash.code_utils.config_consts import TOTAL_LOOPING_TIME
+from codeflash.code_utils.config_consts import get_total_looping_time
 from codeflash.code_utils.coverage_utils import prepare_coverage_files
 from codeflash.models.models import TestFiles, TestType
 
@@ -37,7 +37,7 @@ def run_behavioral_tests(
     pytest_timeout: int | None = None,
     pytest_cmd: str = "pytest",
     verbose: bool = False,
-    pytest_target_runtime_seconds: int = TOTAL_LOOPING_TIME,
+    pytest_target_runtime_seconds: int = get_total_looping_time(),
     enable_coverage: bool = False,
 ) -> tuple[Path, subprocess.CompletedProcess, Path | None, Path | None]:
     if test_framework == "pytest":
@@ -151,7 +151,7 @@ def run_line_profile_tests(
     cwd: Path,
     test_framework: str,
     *,
-    pytest_target_runtime_seconds: float = TOTAL_LOOPING_TIME,
+    pytest_target_runtime_seconds: float = get_total_looping_time(),
     verbose: bool = False,
     pytest_timeout: int | None = None,
     pytest_min_loops: int = 5,  # noqa: ARG001
@@ -237,7 +237,7 @@ def run_benchmarking_tests(
     cwd: Path,
     test_framework: str,
     *,
-    pytest_target_runtime_seconds: float = TOTAL_LOOPING_TIME,
+    pytest_target_runtime_seconds: float = get_total_looping_time(),
     verbose: bool = False,
     pytest_timeout: int | None = None,
     pytest_min_loops: int = 5,